Related Articles
Count of subarrays whose sum is a perfect square
• Difficulty Level : Medium
• Last Updated : 06 Nov, 2020

Given an array arr[] with positive and negative elements, the task is to count all subarrays whose sum is a perfect square.
Examples:

Input: arr[] = {2, 3, -5, 6, -7, 4};
Output:
Explanation:
Subarrays {2, 3, -5}, {-5, 6}, {3, -5, 6}, {3, -5, 6, -7, 4} and {4} with sum is 0, 1, 4, 1 and 4 respectively have perfect square sum.
Input: arr[] = {3, -6, 4, -2, 7};
Output:
Explanation: {3, -6, 4}, {4}, {4, -2, 7} are the subarrays with perfect square sum.

Naive Approach:
A simple solution would be to generate all possible subarrays. While traversing, keep track of the subarray sum. Keep a count of all subarrays whose sum is a perfect square.
Efficient Solution: The idea is to use a prefix sum array to solve the given problem. • Create a prefixSum array and store it’s prefix sum.
• Traverse the prefixSum array and identify it’s minimum value i.e (prefixMin).
• Now, create an unordered map which can be used to store frequency of current prefixSum, while traversing the prefixSum array.
• Initialize the 0th key-index of the map with value 1, as 0 is a perfect square.
• Traverse the prefixSum array with a nested loop.
• For each prefixSum element, the nested loop is going to find the mapKey = (prefixSum[i] – j*j), if available in the map index.
• If (prefixSum[i] – j*j) is already available in the map, we update our counter with the index value of (prefixSum[i] – j*j).
• The idea is to check the current prefixSum value with all the squares (j*j) till the difference reaches prefixMin.
• Now, increment the map with index of the current prefixSum by 1 with every iteration of the outer loop.
• The underlying concept is that we keep searching from (prefixSum[i] – j*j ) because, if one part is the array is (prefixSum[i] – j*j ), then the other part of the array would be (j*j) i.e a perfect square sum.
• You can see in the above diagram that the totalSum is actually the prefixSum, which is used for that purpose.

Below is the implementation of the above approach:

## C++

 `// C++ code for the above approach.``#include ``using` `namespace` `std;` `#define lli long long int` `// Function to find count of subarrays``// whose sum is a perfect square.``lli countSubarrays(``int` `arr[],``                   ``int` `n)``{``    ``// to search for index with``    ``// (current prefix sum - j*j)``    ``unordered_map<``int``, ``int``> mp;` `    ``// storing the prefix sum``    ``int` `prefixSum[n];` `    ``// used to track the minimum``    ``// value in prefixSum``    ``int` `prefixMin = 0;` `    ``prefixSum = arr;``    ``prefixMin = min(prefixMin,``                    ``prefixSum);` `    ``// Calculating the prefixSum``    ``// and tracking the prefixMin``    ``for` `(``int` `i = 1; i < n; i++) {` `        ``prefixSum[i] = prefixSum[i - 1]``                       ``+ arr[i];` `        ``// below statement is used if``        ``// array contains``        ``// negative numbers``        ``prefixMin = min(prefixMin,``                        ``prefixSum[i]);``    ``}` `    ``// counts the no of subarrays``    ``// with perfect square sum``    ``lli countSubs = 0;` `    ``// as 0 is a perfect square,``    ``// so we initialize 0th``    ``// index-key with value 1``    ``mp = 1;` `    ``// Here we count the perfect``    ``// square subarray sum by``    ``// searching if there is a``    ``// prefix with``    ``// sum = (current prefixSum - (sq*sq))``    ``for` `(``int` `i = 0; i < n; i++) {``        ``for` `(``int` `j = 0;``             ``prefixSum[i] - j * j >= prefixMin;``             ``j++) {` `            ``if` `(mp.find(prefixSum[i] - j * j)``                ``!= mp.end())` `                ``// increasing our subarray count``                ``countSubs += mp[prefixSum[i]``                                ``- j * j];``        ``}` `        ``// increasing the current prefixSum``        ``// index value in map by 1 to count``        ``// the other perfect squares while``        ``// traversing further``        ``mp[prefixSum[i]]++;``    ``}` `    ``return` `countSubs;``}` `// Driver code``int` `main()``{``    ``int` `arr[] = { 2, 3, -5,``                  ``6, -7, 4 };``    ``int` `n = ``sizeof``(arr) / ``sizeof``(arr);` `    ``lli ans = countSubarrays(arr, n);` `    ``// printing the result``    ``cout << ans;` `    ``return` `0;``}`

## Java

 `// Java code for``// the above approach.``import` `java.util.*;``class` `GFG{``  ` `// Function to find count of``// subarrays whose sum is``// a perfect square.``static` `long` `countSubarrays(``int` `arr[],``                           ``int` `n)``{``  ``// To search for index with``  ``// (current prefix sum - j*j)``  ``HashMap mp = ``new` `HashMap();` `  ``// Storing the prefix sum``  ``int` `[]prefixSum = ``new` `int``[n];` `  ``// Used to track the minimum``  ``// value in prefixSum``  ``int` `prefixMin = ``0``;` `  ``prefixSum[``0``] = arr[``0``];``  ``prefixMin = Math.min(prefixMin,``                       ``prefixSum[``0``]);` `  ``// Calculating the prefixSum``  ``// and tracking the prefixMin``  ``for` `(``int` `i = ``1``; i < n; i++)``  ``{``    ``prefixSum[i] = prefixSum[i - ``1``] + arr[i];` `    ``// Below statement is used if``    ``// array contains``    ``// negative numbers``    ``prefixMin = Math.min(prefixMin,``                         ``prefixSum[i]);``  ``}` `  ``// Counts the no of subarrays``  ``// with perfect square sum``  ``long` `countSubs = ``0``;` `  ``// As 0 is a perfect square,``  ``// so we initialize 0th``  ``// index-key with value 1``  ``mp.put(``0``, ``1``);` `  ``// Here we count the perfect``  ``// square subarray sum by``  ``// searching if there is a``  ``// prefix with``  ``// sum = (current prefixSum - (sq*sq))``  ``for` `(``int` `i = ``0``; i < n; i++)``  ``{``    ``for` `(``int` `j = ``0``;``             ``prefixSum[i] - j *``             ``j >= prefixMin; j++)``    ``{``      ``if` `(mp.containsKey(prefixSum[i] - j * j))` `        ``// Increasing our subarray count``        ``countSubs += mp.get(prefixSum[i] -``                            ``j * j);``    ``}` `    ``// Increasing the current prefixSum``    ``// index value in map by 1 to count``    ``// the other perfect squares while``    ``// traversing further``    ``if``(mp.containsKey(prefixSum[i]))``    ``{``      ``mp.put(prefixSum[i],``      ``mp.get(prefixSum[i]) + ``1``);``    ``}``    ``else``    ``{``      ``mp.put(prefixSum[i], ``1``);``    ``}``  ``}` `  ``return` `countSubs;``}` `// Driver code``public` `static` `void` `main(String[] args)``{``  ``int` `arr[] = {``2``, ``3``, -``5``,``               ``6``, -``7``, ``4``};``  ``int` `n = arr.length;``  ``long` `ans = countSubarrays(arr, n);` `  ``// Printing the result``  ``System.out.print(ans);``}``}` `// This code is contributed by Princi Singh`

## Python3

 `# Python3 code for the above approach.``from` `collections ``import` `defaultdict` `# Function to find count of subarrays``# whose sum is a perfect square.``def` `countSubarrays(arr, n):``    ` `    ``# To search for index with``    ``# (current prefix sum - j*j)``    ``mp ``=` `defaultdict(``lambda``:``0``)``    ` `    ``# Storing the prefix sum``    ``prefixSum ``=` `[``0``] ``*` `n``    ` `    ``# Used to track the minimum``    ``# value in prefixSum``    ``prefixMin ``=` `0``    ` `    ``prefixSum[``0``] ``=` `arr[``0``]``    ``prefixMin ``=` `min``(prefixMin, prefixSum[``0``])``    ` `    ``# Calculating the prefixSum``    ``# and tracking the prefixMin``    ``for` `i ``in` `range``(``1``, n):``        ``prefixSum[i] ``=` `prefixSum[i ``-` `1``] ``+` `arr[i]``        ` `        ``# Below statement is used if``        ``# array contains negative numbers``        ``prefixMin ``=` `min``(prefixMin, prefixSum[i])``        ` `    ``# Counts the no of subarrays``    ``# with perfect square sum``    ``countSubs ``=` `0``    ` `    ``# As 0 is a perfect square,``    ``# so we initialize 0th``    ``# index-key with value 1``    ``mp[``0``] ``=` `1``    ` `    ``# Here we count the perfect``    ``# square subarray sum by``    ``# searching if there is a``    ``# prefix with``    ``# sum = (current prefixSum - (sq*sq))``    ``for` `i ``in` `range``(n):``        ``j ``=` `0``        ` `        ``while` `prefixSum[i] ``-` `j ``*` `j >``=` `prefixMin:``            ``if` `prefixSum[i] ``-` `j ``*` `j ``in` `mp:``                ` `                ``# Increasing our subarray count``                ``countSubs ``+``=` `mp[prefixSum[i] ``-` `j ``*` `j]``            ``j ``+``=` `1``            ` `        ``# Increasing the current prefixSum``        ``# index value in map by 1 to count``        ``# the other perfect squares while``        ``# traversing further``        ``mp[prefixSum[i]] ``+``=` `1``        ` `    ``return` `countSubs` `# Driver code``arr ``=` `[ ``2``, ``3``, ``-``5``, ``6``, ``-``7``, ``4` `]``n ``=` `len``(arr)``ans ``=` `countSubarrays(arr, n)``    ` `# Printing the result``print``(ans)` `# This code is contributed by Shivam Singh`

## C#

 `// C# code for``// the above approach.``using` `System;``using` `System.Collections.Generic;``class` `GFG{``  ` `// Function to find count of``// subarrays whose sum is``// a perfect square.``static` `long` `countSubarrays(``int` `[]arr,``                           ``int` `n)``{``  ``// To search for index with``  ``// (current prefix sum - j*j)``  ``Dictionary<``int``,``             ``int``> mp =``             ``new` `Dictionary<``int``,``                            ``int``>();` `  ``// Storing the prefix sum``  ``int` `[]prefixSum = ``new` `int``[n];` `  ``// Used to track the minimum``  ``// value in prefixSum``  ``int` `prefixMin = 0;` `  ``prefixSum = arr;``  ``prefixMin = Math.Min(prefixMin,``                       ``prefixSum);` `  ``// Calculating the prefixSum``  ``// and tracking the prefixMin``  ``for` `(``int` `i = 1; i < n; i++)``  ``{``    ``prefixSum[i] = prefixSum[i - 1] +``                   ``arr[i];` `    ``// Below statement is used if``    ``// array contains``    ``// negative numbers``    ``prefixMin = Math.Min(prefixMin,``                         ``prefixSum[i]);``  ``}` `  ``// Counts the no of subarrays``  ``// with perfect square sum``  ``long` `countSubs = 0;` `  ``// As 0 is a perfect square,``  ``// so we initialize 0th``  ``// index-key with value 1``  ``mp.Add(0, 1);` `  ``// Here we count the perfect``  ``// square subarray sum by``  ``// searching if there is a``  ``// prefix with``  ``// sum = (current prefixSum -``  ``// (sq*sq))``  ``for` `(``int` `i = 0; i < n; i++)``  ``{``    ``for` `(``int` `j = 0; prefixSum[i] - j *``             ``j >= prefixMin; j++)``    ``{``      ``if` `(mp.ContainsKey(prefixSum[i] -``                         ``j * j))` `        ``// Increasing our subarray count``        ``countSubs += mp[prefixSum[i] -``                        ``j * j];``    ``}` `    ``// Increasing the current prefixSum``    ``// index value in map by 1 to count``    ``// the other perfect squares while``    ``// traversing further``    ``if``(mp.ContainsKey(prefixSum[i]))``    ``{``      ``mp[prefixSum[i]]++;   ``    ``}``    ``else``    ``{``      ``mp.Add(prefixSum[i], 1);``    ``}``  ``}` `  ``return` `countSubs;``}` `// Driver code``public` `static` `void` `Main(String[] args)``{``  ``int` `[]arr = {2, 3, -5,``               ``6, -7, 4};``  ``int` `n = arr.Length;``  ``long` `ans = countSubarrays(arr, n);` `  ``// Printing the result``  ``Console.Write(ans);``}``}` `// This code is contributed by gauravrajput1`
Output:
```5

```

Time Complexity: O(N * sqrt(K))
Auxiliary Space: O(N)

Attention reader! Don’t stop learning now. Get hold of all the important mathematical concepts for competitive programming with the Essential Maths for CP Course at a student-friendly price. To complete your preparation from learning a language to DS Algo and many more,  please refer Complete Interview Preparation Course.

My Personal Notes arrow_drop_up