Count of subarrays whose sum is a perfect square - GeeksforGeeks

24 Nov, 2021
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: 5 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: 3 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++
Java
Python3
C#
Javascript
// C++ code for the above approach.#include <bits/stdc++.h>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 codeint 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 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<Integer,          Integer> mp = new HashMap<Integer,                                    Integer>();   // 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 codepublic 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 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 codearr = [ 2, 3, -5, 6, -7, 4 ]n = len(arr)ans = countSubarrays(arr, n)     # Printing the resultprint(ans) # This code is contributed by Shivam Singh
// 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 codepublic 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
<script> // Javascript code for// the above approach. // Function to find count of// subarrays whose sum is// a perfect square.function countSubarrays(arr, n){  // To search for index with  // (current prefix sum - j*j)  let mp = new Map();    // Storing the prefix sum  let prefixSum = Array.from({length: n}, (_, i) => 0);    // Used to track the minimum  // value in prefixSum  let prefixMin = 0;    prefixSum[0] = arr[0];  prefixMin = Math.min(prefixMin,                       prefixSum[0]);    // Calculating the prefixSum  // and tracking the prefixMin  for (let 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  let countSubs = 0;    // As 0 is a perfect square,  // so we initialize 0th  // index-key with value 1  mp.set(0, 1);    // Here we count the perfect  // square subarray sum by  // searching if there is a  // prefix with  // sum = (current prefixSum - (sq*sq))  for (let i = 0; i < n; i++)  {    for (let j = 0;             prefixSum[i] - j *             j >= prefixMin; j++)    {      if (mp.has(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.has(prefixSum[i]))    {      mp.set(prefixSum[i],      mp.get(prefixSum[i]) + 1);    }    else    {      mp.set(prefixSum[i], 1);    }  }    return countSubs;} // Driver code           let arr = [2, 3, -5,               6, -7, 4];  let n = arr.length;  let ans = countSubarrays(arr, n);    // Printing the result  document.write(ans); // This code is contributed by souravghosh0416.</script>
5
 
Time Complexity: O(N * sqrt(K)) Auxiliary Space: O(N)
SHIVAMSINGH67
nidhi_biet
princi singh
GauravRajput1
souravghosh0416
simranarora5sos
maths-perfect-square
subarray-sum
Arrays
Hash
Mathematical
Arrays
Hash
Mathematical
Writing code in comment?
Please use ide.geeksforgeeks.org,
generate link and share the link here.
Stack Data Structure (Introduction and Program)
Top 50 Array Coding Problems for Interviews
Multidimensional Arrays in Java
Introduction to Arrays
Linear Search
Given an array A[] and a number x, check for pair in A[] with sum as x (aka Two Sum)
Internal Working of HashMap in Java
Hashing | Set 1 (Introduction)
Hashing | Set 3 (Open Addressing)
Hashing | Set 2 (Separate Chaining)