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: 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 repectively have perfect square sum.

Input: arr[] = {3, -6, 4, -2, 7};
Output: 3
Exxplanation: {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 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++

filter_none

edit
close

play_arrow

link
brightness_4
code

// 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 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;
}

chevron_right


Python3

filter_none

edit
close

play_arrow

link
brightness_4
code

# 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

chevron_right


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 DSA concepts with the DSA Self Paced Course at a student-friendly price and become industry ready.




My Personal Notes arrow_drop_up

Check out this Author's contributed articles.

If you like GeeksforGeeks and would like to contribute, you can also write an article using contribute.geeksforgeeks.org or mail your article to contribute@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.

Please Improve this article if you find anything incorrect by clicking on the "Improve Article" button below.



Improved By : SHIVAMSINGH67, nidhi_biet