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++
// 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
// 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
# 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#
// 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 |
Javascript
<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> |
Output:
5
Time Complexity: O(N * sqrt(K))
Auxiliary Space: O(N)
Feeling lost in the world of random DSA topics, wasting time without progress? It’s time for a change! Join our DSA course, where we’ll guide you on an exciting journey to master DSA efficiently and on schedule.
Ready to dive in? Explore our Free Demo Content and join our DSA course, trusted by over 100,000 neveropen!
Ready to dive in? Explore our Free Demo Content and join our DSA course, trusted by over 100,000 neveropen!
