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:

- Number of subarrays having absolute sum greater than K | Set-2
- Number of subarrays having sum exactly equal to k
- Number of subarrays having sum less than K
- Number of subarrays having sum in a given range
- Find an array of size N having exactly K subarrays with sum S
- Count of subarrays having sum as a perfect cube
- Count subarrays having sum of elements at even and odd positions equal
- Count of ways to split an Array into three contiguous Subarrays having increasing Sum
- Count of subarrays having sum equal to its length
- Number of subarrays having product less than K
- Number of ways to select equal sized subarrays from two arrays having atleast K equal pairs of elements
- Split array into minimum number of subarrays having GCD of its first and last element exceeding 1
- Ways to form an array having integers in given range such that total sum is divisible by 2
- Count of subarrays having exactly K distinct elements
- Count of subarrays having exactly K prime numbers
- Count of subarrays of an Array having all unique digits
- Split an Array to maximize subarrays having equal count of odd and even elements for a cost not exceeding K
- Count of subarrays of size K with elements having even frequencies
- Count of subarrays of size K having at least one pair with absolute difference divisible by K-1
- Longest Subarrays having each Array element as the maximum

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.