Open In App

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

Improve
Improve
Like Article
Like
Save
Share
Report

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

Input: arr[] = { 2, 2, 2, 2 } K = 2 
Output:
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:

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:  

C++




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


Java




// Java implementation of the
// above approach
import java.util.*;
class GFG{
 
static final int MAX = 100005;
   
// partial_sum
static long[] partial_sum(long []prefix_sum,
                          int[]arr, int n)
{
  for (int i = 1; i <= n; i++)
  {
    prefix_sum[i] = (prefix_sum[i - 1] +
                     arr[i - 1]);
  }
 
  return prefix_sum;
}
   
// Function to count number of
// sub-arrays whose sum is k^p
// where p>=0
static int countSubarrays(int []arr,
                          int n, int k)
{
  long []prefix_sum = new long[MAX];
  prefix_sum[0] = 0;
  prefix_sum = partial_sum(prefix_sum ,
                           arr, n);
  int sum;
 
  if (k == 1)
  {
    sum = 0;
    HashMap<Long,
            Integer> m = new HashMap<>();
 
    for (int i = n; i >= 0; i--)
    {
      // If m[a+b] = c, then add c to
      // the current sum.
      if (m.containsKey(prefix_sum[i] + 1))
        sum += m.get(prefix_sum[i] + 1);
 
      // Increase count of prefix sum.
      if(m.containsKey(prefix_sum[i]))
        m.put(prefix_sum[i],
        m.get(prefix_sum[i]) + 1);
      else
        m.put(prefix_sum[i], 1);
    }
    return sum;
  }
 
  if (k == -1)
  {
    sum = 0;
    HashMap<Long,
            Integer> m = new HashMap<>();
 
    for (int i = n; i >= 0; i--)
    {
      // If m[a+b] = c, then add c to
      // the current sum.
      if (m.containsKey(prefix_sum[i] + 1))
        sum += m.get(prefix_sum[i] + 1);
 
      if (m.containsKey(prefix_sum[i] - 1))
        sum += m.get(prefix_sum[i] - 1);
 
      // Increase count of prefix sum.
      if(m.containsKey(prefix_sum[i]))
        m.put(prefix_sum[i],
        m.get(prefix_sum[i]) + 1);
      else
        m.put(prefix_sum[i], 1);
    }
    return sum;
  }
 
  sum = 0;
 
  // b = k^p, p>=0
  long b, l = 100000000000000L;
  HashMap<Long,
          Integer> m = new HashMap<>();
 
  for (int i = n; i >= 0; i--)
  {
    b = 1;
    while (true)
    {
      // k^m can be maximum equal
      // to 10^14.
      if (b > l)
        break;
 
      // If m[a+b] = c, then add c to
      // the current sum.
      if (m.containsKey(prefix_sum[i] + b))
        sum += m.get(prefix_sum[i] + b);
 
      b *= k;
    }
 
    // Increase count of prefix sum.
    if(m.containsKey(prefix_sum[i]))
      m.put((prefix_sum[i]),
      m.get(prefix_sum[i]) + 1);
    else
      m.put((prefix_sum[i]), 1);
  }
  return sum;
}
 
// Driver code
public static void main(String[] args)
{
  int arr[] = {2, 2, 2, 2};
  int n = arr.length;
  int k = 2;
  System.out.print(countSubarrays(arr,
                                  n, k));
}
}
 
// This code is contributed by Rajput-Ji


Python3




# Python3 implementation of
# the above approach
from collections import defaultdict
MAX = 100005
 
def partial_sum(prefix_sum,
                arr, n):
   
  for i in range(1 , n + 1):
    prefix_sum[i] = (prefix_sum[i - 1] +
                     arr[i - 1])
  return prefix_sum
     
# Function to count number of
# sub-arrays whose sum is k^p
# where p>=0
def countSubarrays(arr, n, k):
 
    prefix_sum = [0] * MAX
    prefix_sum[0] = 0
      
    prefix_sum = partial_sum(prefix_sum,
                             arr, n)
    if (k == 1):
        sum = 0
        m = defaultdict(int)
 
        for i in range(n, -1, -1):
 
            # If m[a+b] = c, then add
            # c to the current sum.
            if ((prefix_sum[i] + 1) in m):
                sum += m[prefix_sum[i] + 1]
 
            # Increase count of prefix sum.
            m[prefix_sum[i]] += 1
 
        return sum
 
    if (k == -1):
        sum = 0
        m = defaultdict(int)
 
        for i in range(n, -1, -1):
 
            # If m[a+b] = c, then add c
            # to the current sum.
            if ((prefix_sum[i] + 1) in m):
                sum += m[prefix_sum[i] + 1]
 
            if ((prefix_sum[i] - 1) in m):
                sum += m[prefix_sum[i] - 1]
 
            # Increase count of prefix sum.
            m[prefix_sum[i]] += 1
 
        return sum
 
    sum = 0
 
    # b = k^p, p>=0
    m = defaultdict(int)
 
    for i in range(n, -1, -1):
        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 ((prefix_sum[i] + b) in m):
                sum += m[prefix_sum[i] + b]
 
            b *= k
 
        # Increase count of prefix
        # sum.
        m[prefix_sum[i]] += 1
    return sum
 
# Driver code
if __name__ == "__main__":
 
    arr = [2, 2, 2, 2]
    n = len(arr)
    k = 2
    print(countSubarrays(arr, n, k))
 
# This code is contributed by Chitranayal


C#




// C# implementation of the
// above approach
using System;
using System.Collections.Generic;
class GFG{
 
static readonly int MAX = 100005;
   
// partial_sum
static long[] partial_sum(long []prefix_sum,
                          int[]arr, int n)
{
  for (int i = 1; i <= n; i++)
  {
    prefix_sum[i] = (prefix_sum[i - 1] +
                     arr[i - 1]);
  }
 
  return prefix_sum;
}
   
// Function to count number of
// sub-arrays whose sum is k^p
// where p>=0
static int countSubarrays(int []arr,
                          int n, int k)
{
  long []prefix_sum = new long[MAX];
  prefix_sum[0] = 0;
  prefix_sum = partial_sum(prefix_sum ,
                           arr, n);
  int sum;
 
  if (k == 1)
  {
    sum = 0;
    Dictionary<long,
               int> mp =
               new Dictionary<long,
                              int>();
 
    for (int i = n; i >= 0; i--)
    {
      // If m[a+b] = c, then add c to
      // the current sum.
      if (mp.ContainsKey(prefix_sum[i] + 1))
        sum += mp[prefix_sum[i] + 1];
 
      // Increase count of prefix sum.
      if(mp.ContainsKey(prefix_sum[i]))
        mp.Add(prefix_sum[i],
        mp[prefix_sum[i]] + 1);
      else
        mp.Add(prefix_sum[i], 1);
    }
    return sum;
  }
 
  if (k == -1)
  {
    sum = 0;
    Dictionary<long,
               int> map =
               new Dictionary<long,
                              int>();
 
    for (int i = n; i >= 0; i--)
    {
      // If m[a+b] = c, then add c to
      // the current sum.
      if (map.ContainsKey(prefix_sum[i] + 1))
        sum += map[prefix_sum[i] + 1];
 
      if (map.ContainsKey(prefix_sum[i] - 1))
        sum += map[prefix_sum[i] - 1];
 
      // Increase count of prefix sum.
      if(map.ContainsKey(prefix_sum[i]))
        map.Add(prefix_sum[i],
        map[prefix_sum[i]] + 1);
      else
        map.Add(prefix_sum[i], 1);
    }
    return sum;
  }
 
  sum = 0;
 
  // b = k^p, p>=0
  long b, l = 100000000000000L;
  Dictionary<long,
             int> m =
             new Dictionary<long,
                            int>();
 
  for (int i = n; i >= 0; i--)
  {
    b = 1;
    while (true)
    {
      // k^m can be maximum equal
      // to 10^14.
      if (b > l)
        break;
 
      // If m[a+b] = c, then add c to
      // the current sum.
      if (m.ContainsKey(prefix_sum[i] + b))
        sum += m[prefix_sum[i] + b];
 
      b *= k;
    }
 
    // Increase count of prefix sum.
    if(m.ContainsKey(prefix_sum[i]))
      m.Add((prefix_sum[i]),
      m[prefix_sum[i]] + 1);
    else
      m.Add((prefix_sum[i]), 1);
  }
  return sum;
}
 
// Driver code
public static void Main(String[] args)
{
  int []arr = {2, 2, 2, 2};
  int n = arr.Length;
  int k = 2;
  Console.Write(countSubarrays(arr,
                               n, k));
}
}
 
// This code is contributed by shikhasingrajput


Javascript




<script>
 
// JavaScript implementation of the
// above approach
 
let MAX = 100005;
 
// partial_sum
function partial_sum(prefix_sum,arr,n)
{
    for (let i = 1; i <= n; i++)
  {
    prefix_sum[i] = (prefix_sum[i - 1] +
                     arr[i - 1]);
  }
  
  return prefix_sum;
}
 
// Function to count number of
// sub-arrays whose sum is k^p
// where p>=0
function countSubarrays(arr,n,k)
{
    let prefix_sum = new Array(MAX);
  prefix_sum[0] = 0;
  prefix_sum = partial_sum(prefix_sum ,
                           arr, n);
  let sum;
  
  if (k == 1)
  {
    sum = 0;
    let m = new Map();
  
    for (let i = n; i >= 0; i--)
    {
      // If m[a+b] = c, then add c to
      // the current sum.
      if (m.has(prefix_sum[i] + 1))
        sum += m.get(prefix_sum[i] + 1);
  
      // Increase count of prefix sum.
      if(m.has(prefix_sum[i]))
        m.set(prefix_sum[i],
        m.get(prefix_sum[i]) + 1);
      else
        m.set(prefix_sum[i], 1);
    }
    return sum;
  }
  
  if (k == -1)
  {
    sum = 0;
    let m = new Map();
  
    for (let i = n; i >= 0; i--)
    {
      // If m[a+b] = c, then add c to
      // the current sum.
      if (m.has(prefix_sum[i] + 1))
        sum += m.get(prefix_sum[i] + 1);
  
      if (m.has(prefix_sum[i] - 1))
        sum += m.get(prefix_sum[i] - 1);
  
      // Increase count of prefix sum.
      if(m.has(prefix_sum[i]))
        m.set(prefix_sum[i],
        m.get(prefix_sum[i]) + 1);
      else
        m.set(prefix_sum[i], 1);
    }
    return sum;
  }
  
  sum = 0;
  
  // b = k^p, p>=0
  let b, l = 100000000000000;
  let m = new Map();
  
  for (let i = n; i >= 0; i--)
  {
    b = 1;
    while (true)
    {
      // k^m can be maximum equal
      // to 10^14.
      if (b > l)
        break;
  
      // If m[a+b] = c, then add c to
      // the current sum.
      if (m.has(prefix_sum[i] + b))
        sum += m.get(prefix_sum[i] + b);
  
      b *= k;
    }
  
    // Increase count of prefix sum.
    if(m.has(prefix_sum[i]))
      m.set((prefix_sum[i]),
      m.get(prefix_sum[i]) + 1);
    else
      m.set((prefix_sum[i]), 1);
  }
  return sum;
}
 
// Driver code
let arr=[2, 2, 2, 2];
let n = arr.length;
let k = 2;
document.write(countSubarrays(arr,
                                  n, k));
 
 
// This code is contributed by avanitrachhadiya2155
 
</script>


Output: 

8

 



Last Updated : 31 May, 2021
Like Article
Save Article
Previous
Next
Share your thoughts in the comments
Similar Reads