📜  子数组的总和形式为k ^ m,m> = 0

📅  最后修改于: 2021-04-27 21:03:58             🧑  作者: Mango

给定一个整数k和一个数组arr [] ,任务是计算总和等于k的某个正整数幂的子数组的数量。
例子:

天真的方法:天真的方法是遍历所有子数组,并检查每个子数组的总和是否等于k的某个整数幂。
高效的方法:更好的方法是维护一个前缀和数组prefix_sum和一个映射m ,该映射m将前缀和映射到其计数。 m [a] = 1表示a是某个前缀的前缀和。

注意:k = 1和k = -1需要分开处理。
下面是上述方法的实现:

C++
// C++ implementation of the above approach
#include 
 
#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 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 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 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 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 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 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 mp =
               new Dictionary();
 
    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 map =
               new Dictionary();
 
    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 m =
             new Dictionary();
 
  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


输出:
8