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++ 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;
} |
// 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 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 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# 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 |
<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)