GFG App
Open App
Browser
Continue

# Count of subarrays whose sum is a perfect square

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 that can be used to store the 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[0] = arr[0];` `    ``prefixMin = 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 = 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[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.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[0]);`   `    ``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[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.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`

## Javascript

 ``

Output:

`5`

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

My Personal Notes arrow_drop_up