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

Last Updated : 31 May, 2021

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: 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 i^th 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 = k^p and c is the value fetched from the map, then it means that starting from the i^th 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# 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:

Number of subarrays having product less than K

rohan23chhabra

Improve

Article Tags :

Practice Tags :

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

Similar Reads

Thank You!

What kind of Experience do you want to share?