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 direction and carefully follow the below discussion. Consider that while traversing the array we are at the ith index and after traversing an index we perform the operation op = m[prefix_sum[i]]++. Hence when we are at index i, op hasn’t been performed yet. See the code for a vivid explanation.

If m[a + b] = c where a = prefix_sum[i], b = kp and c is the value fetched from the map, then it means that starting from the ith index to the end of the array, there are c sub-arrays whose sum is equal to b. Add c to the current sum.
This is because for every index j > i, m[prefix_sum[j]]++ has been performed. Therefore the map has information about prefix sums of prefixes ending at j > i. On adding b to the prefix sum we can get the count of all those sums a + b which will indicate that a sub-array exists that has sum equal to b.

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

Below is the implementation of the above approach:

filter_none

edit
close

play_arrow

link
brightness_4
code

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


Output:

8


My Personal Notes arrow_drop_up

An enthusiastic Java and web developer with a little affinity for tea, cricket, English, etymology, and reading

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.