# Number of subarrays having sum of the form k^m, m >= 0

Given an integer **k** and an array **arr[]**, the task is to count the number of sub-arrays which have the sum equal to some positive integral power of k.

**Examples:**

Input:arr[] = { 2, 2, 2, 2 } K = 2

Output:8

Sub-arrays with below indexes are valid:

[1, 1], [2, 2], [3, 3], [4, 4], [1, 2],

[2, 3], [3, 4], [1, 4]

Input:arr[] = { 3, -6, -3, 12 } K = -3

Output:3

**Naive Approach**: A naive approach is to traverse through all the sub-arrays and check for every sub-array whether its sum is equal to some integral power of k.

**Efficient Approach**: A better approach is to maintain a prefix sum array **prefix_sum** and a map **m** which maps a prefix sum to its count. *m[a] = 1 means that a is a prefix sum of some prefix.*

Iterate through the array in the

backward directionand carefully follow the below discussion. Consider that while traversing the array we are at the i^{th}index and after traversing an index we perform the operationop = m[prefix_sum[i]]++. Hence when we are at indexi,ophasn’t been performed yet. See the code for a vivid explanation.If

m[a + b] = cwherea = prefix_sum[i], b = k, then it means that starting from the i^{p}and c is the value fetched from the map^{th}index to the end of the array, there are c sub-arrays whose sum is equal to b. Addcto the current sum.

This is because for every indexj > i, m[prefix_sum[j]]++has been performed. Therefore the map has information about prefix sums of prefixes ending atj > i. On addingbto the prefix sum we can get the count of all those sumsa + bwhich will indicate that a sub-array exists that has sum equal tob.

**Note**: k = 1 and k = -1 need to be handled separately.

Below is the implementation of the above approach:

`// C++ implementation of the above approach ` `#include <bits/stdc++.h> ` ` ` `#define ll long long ` `#define MAX 100005 ` ` ` `using` `namespace` `std; ` ` ` `// Function to count number of sub-arrays ` `// whose sum is k^p where p>=0 ` `ll countSubarrays(` `int` `* arr, ` `int` `n, ` `int` `k) ` `{ ` ` ` `ll prefix_sum[MAX]; ` ` ` `prefix_sum[0] = 0; ` ` ` ` ` `partial_sum(arr, arr + n, prefix_sum + 1); ` ` ` ` ` `ll sum; ` ` ` ` ` `if` `(k == 1) { ` ` ` ` ` `sum = 0; ` ` ` `map<ll, ` `int` `> m; ` ` ` ` ` `for` `(` `int` `i = n; i >= 0; i--) { ` ` ` ` ` `// If m[a+b] = c, then add c to the current sum. ` ` ` `if` `(m.find(prefix_sum[i] + 1) != m.end()) ` ` ` `sum += m[prefix_sum[i] + 1]; ` ` ` ` ` `// Increase count of prefix sum. ` ` ` `m[prefix_sum[i]]++; ` ` ` `} ` ` ` ` ` `return` `sum; ` ` ` `} ` ` ` ` ` `if` `(k == -1) { ` ` ` ` ` `sum = 0; ` ` ` `map<ll, ` `int` `> m; ` ` ` ` ` `for` `(` `int` `i = n; i >= 0; i--) { ` ` ` ` ` `// If m[a+b] = c, then add c to the current sum. ` ` ` `if` `(m.find(prefix_sum[i] + 1) != m.end()) ` ` ` `sum += m[prefix_sum[i] + 1]; ` ` ` ` ` `if` `(m.find(prefix_sum[i] - 1) != m.end()) ` ` ` `sum += m[prefix_sum[i] - 1]; ` ` ` ` ` `// Increase count of prefix sum. ` ` ` `m[prefix_sum[i]]++; ` ` ` `} ` ` ` ` ` `return` `sum; ` ` ` `} ` ` ` ` ` `sum = 0; ` ` ` ` ` `// b = k^p, p>=0 ` ` ` `ll b; ` ` ` `map<ll, ` `int` `> m; ` ` ` ` ` `for` `(` `int` `i = n; i >= 0; i--) { ` ` ` ` ` `b = 1; ` ` ` `while` `(` `true` `) { ` ` ` ` ` `// k^m can be maximum equal to 10^14. ` ` ` `if` `(b > 100000000000000) ` ` ` `break` `; ` ` ` ` ` `// If m[a+b] = c, then add c to the current sum. ` ` ` `if` `(m.find(prefix_sum[i] + b) != m.end()) ` ` ` `sum += m[prefix_sum[i] + b]; ` ` ` ` ` `b *= k; ` ` ` `} ` ` ` ` ` `// Increase count of prefix sum. ` ` ` `m[prefix_sum[i]]++; ` ` ` `} ` ` ` ` ` `return` `sum; ` `} ` ` ` `// Driver code ` `int` `main() ` `{ ` ` ` `int` `arr[] = { 2, 2, 2, 2 }; ` ` ` `int` `n = ` `sizeof` `(arr) / ` `sizeof` `(arr[0]); ` ` ` `int` `k = 2; ` ` ` ` ` `cout << countSubarrays(arr, n, k); ` ` ` ` ` `return` `0; ` `} ` |

*chevron_right*

*filter_none*

**Output:**

8

## Recommended Posts:

- Differences between number of increasing subarrays and decreasing subarrays in k sized windows
- Number of subarrays having sum less than K
- Find number of subarrays with even sum
- Number of subarrays with given product
- Number of subarrays with m odd numbers
- Number of subarrays have bitwise OR >= K
- Number of subarrays having product less than K
- Maximize the number of subarrays with XOR as zero
- Count the number of subarrays having a given XOR
- Number of subarrays having sum exactly equal to k
- Number of subarrays having sum in a given range
- Number of subarrays for which product and sum are equal
- Count subarrays with equal number of 1's and 0's
- Count the number of non-increasing subarrays
- Number of subarrays having absolute sum greater than K | Set-2

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.