Find Number of Subarrays Having Sum of the Form k * m - m + 0 Using C++

In this article, we will explain everything about solving the number of subarrays having the sum of the form k^m, m >= 0 in C++. Given an array arr[] and an integer K, we need to find the number of subarrays having sum in the form of K^m where m is greater than equal to zero, or we can say we need to find the number of subarrays having sum equal to some non-negative power of K.

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

There are mainly two approaches that come to mind −

Brute Force

In this approach, we will go through all the subarrays and check if they are some positive integral power of K or not; if yes, then we increment the count.

Example

#include <bits/stdc++.h>
#define MAX 1000000
using namespace std;
int main(){
   int arr[] = {2, 2, 2, 2}; // given array
   int k = 2; // given integer
   int n = sizeof(arr) / sizeof(arr[0]); // the size of our array
   int answer = 0; // counter variable
   for(int i = 0; i < n; i++){
      int sum = 0;
      for(int j = i; j < n; j++){ // this will loop will make all the subarrays
         sum += arr[j];
         int b = 1;
         while(b < MAX && sum > b) // k^m Max should be 10^6
            b *= k;
         if(b == sum) // if b == sum then increment count
            answer++;
      }
   }
   cout << answer << "\n";
}

Output

8

However, this approach is not very good as the time complexity of this program is O(N*N*log(K)), where N is the size of our array and K is the integer given by the user.

This complexity is not good as this complexity can be used for higher constraints as it will take too much time to process if the constraints are big, so we will try another approach so that we can use the program for higher constraints.

Efficient Approach

In this approach, we will use a prefix sum and map to reduce our processing that will drastically decrease our time complexity.

Example

#include <bits/stdc++.h>
#define ll long long
#define MAX 1000000
using namespace std;
int main(){
   int arr[] = {2, 2, 2, 2}; // The given array
   int n = sizeof(arr) / sizeof(arr[0]); // size of our array
   int k = 2; // given integer
   ll prefix_sum[MAX];
   prefix_sum[0] = 0;
   partial_sum(arr, arr + n, prefix_sum + 1); // making prefix sum array
   ll sum;
   if (k == 1){
   // we are going to check separately for 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]]++;
      }
      cout << sum << "\n";
   }
   else if (k == -1){
      // we are going to check separately for -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]]++;
      }
      cout << sum << "\n";
   }
   else{
      sum = 0;
      ll b;
      map<ll, int> m;
      for (int i = n; i >= 0; i--){
         b = 1;
         while (b < MAX){ // we are not going to check for more than 10^6
            // 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;
         }
         m[prefix_sum[i]]++;
      }
      cout << sum << "\n";
   }
   return 0;
}

Output

8

Conclusion

We solve a problem to find the Number of subarrays having sum in the form of k^m where m >= 0 in O(nlog(k)log(n)) time complexity. We also learned the C++ program for this problem and the complete approach ( Normal and efficient ) by which we solved this problem. We can write the same program in other languages such as C, java, python, and other languages. Hope you find this article helpful.