Count of Subarray with B times Sum equal to C times Length

Given, an array A[] of N integers, and also given two positive integers B and C, the task is to find the number of subarrays, such that B* Sum of elements of the subarray  = C * Length of the subarray (here * describes multiplication).

Examples:

Input: N = 3, A[] = {3, -1, 1}, B = 1, C = 1
Output: 3
Explanation: The subarrays satisfying the conditions are [3, – 1], [3, – 1, 1], and [1].

• For the subarray [3, -1], sum of the subarray is equal to 2 and length of the subarray is 2. So the condition holds good for it i.e., (1 * 2 = 1 * 2).
• For the subarray [3, -1, 1], sum of the subarray is equal to 3 and length of the subarray is 3. So the condition holds good for it i.e., (1 * 3 = 1 * 3).
• For the subarray [1], sum of the subarray is equal to 1 and length of the subarray is 1. So the condition holds good for it i.e., (1 * 1 = 1 * 1).   

Input: N = 3, A[] = {1, 2, 3}, B = 1, C = 2
Output: 2
Explanation: The subarrays satisfying the conditions are [1, 2, 3] and [2].

Approach: To solve the problem follow the below idea:

The approach is using a map to hash value, and using prefix sum for quick sum calculation of ranges. For example, if a subarray from [l, r] satisfies the conditions, then according to the question,  

• B * (prefix[r] – prefix[l – 1]) = C * (r – l + 1). On Rearranging,
• B * prefix[r] – C * r = B * prefix[l – 1] – C * l + C, so we have all the r terms come on the left side and l terms on the right side So, we can use a map to hash the value on the right side and then search for the value of left side as key in the map.

Steps that were to follow the above approach:

• Create a prefix array to store the prefix sums of array A.
• Create a Map to store the value of the right side as a key, So that we can search for the value of the left side in the Map.
• Iterate over the array and store the right side value in the map and check if the left side value is already present in the Map, if so increment the ans variable that many times.
• Return the answer.

Below is the code to implement the above steps:

C++

// C++ code for the above approach
#include <bits/stdc++.h>
using namespace std;

// This function returns the count of all
// the subarrays satisfying the conditions
int countSubarrays(int N, int A[], int B, int C)
{

    // Declaring the prefix vector to
    // store prefix sums of the array.
    vector<long long> prefix(N + 1);

    // Computing the prefix sums
    for (int i = 1; i <= N; i++) {
        prefix[i] = prefix[i - 1] + A[i - 1];
    }

    // ans will store the count of subarrays
    long long ans = 0;
    map<long long, int> m;

    // Hashing the value on the right side
    // and then search for the value of
    // left side as key in the map
    for (int l = 1; l <= N; l++) {
        m[B * prefix[l - 1] - C * l + C] += 1;
        ans += m[B * prefix[l] - C * l];
    }

    // Returning the result
    return ans;
}

// Driver's code
int main()
{
    int N = 3;
    int A[] = { 1, 2, 3 };
    int B = 1, C = 2;

    // function call
    cout << countSubarrays(N, A, B, C);

    return 0;
}

Java

// Java code for the above approach

import java.io.*;
import java.util.*;

class GFG {

    static int countSubarrays(int N, int[] A, int B, int C)
    {
        // declaring the prefix array to store prefix sums
        // of the array
        int[] prefix = new int[N + 1];

        // computing the prefix sums
        for (int i = 1; i <= N; i++) {
            prefix[i] = prefix[i - 1] + A[i - 1];
        }

        // ans will store the count of subarrays
        int ans = 0;
        Map<Integer, Integer> m = new HashMap<>();

        // Hashing the value on the right side and then
        // search for the value of left side as key in the
        // map
        for (int l = 1; l <= N; l++) {
            m.put(B * prefix[l - 1] - C * l + C,
                  m.getOrDefault(
                      B * prefix[l - 1] - C * l + C, 0)
                      + 1);
            ans += m.getOrDefault(B * prefix[l] - C * l, 0);
        }

        // returning the result
        return ans;
    }

    public static void main(String[] args)
    {
        int N = 3;
        int[] A = { 1, 2, 3 };
        int B = 1, C = 2;

        // Function call
        System.out.print(countSubarrays(N, A, B, C));
    }
}

// This code is contributed by sankar.

Python3

# This function returns the count of all
# the subarrays satisfying the conditions
def countSubarrays(N, A, B, C):
  # Declaring the prefix vector to
  # store prefix sums of the array.
    prefix = [0] * (N + 1)

  # Computing the prefix sums
    for i in range(1, N + 1):
        prefix[i] = prefix[i - 1] + A[i - 1]
  # ans will store the count of subarrays
    ans = 0
    m = {}

  # Hashing the value on the right side
  # and then search for the value of
  # left side as key in the map
    for l in range(1, N + 1):
        m[B * prefix[l - 1] - C * l + C] = m.get(B * prefix[l - 1] - C * l + C, 0) + 1
        ans += m.get(B * prefix[l] - C * l, 0)
  # Returning the result
    return ans
# Driver code
N = 3
A = [1, 2, 3]
B = 1
C = 2

print(countSubarrays(N, A, B, C))

C#

// c# code for the above approach

using System;
using System.Collections.Generic;

public class GFG
{
    static int CountSubarrays(int N, int[] A, int B, int C)
    {
        // declaring the prefix array to store prefix sums
        int[] prefix = new int[N + 1];

        // computing the prefix sums
        for (int i = 1; i <= N; i++)
        {
            prefix[i] = prefix[i - 1] + A[i - 1];
        }

        // ans will store the count of subarrays
        int ans = 0;
        Dictionary<int, int> m = new Dictionary<int, int>();

        // Hashing the value on the right side and then
        // search for the value of left side as key in the map
        for (int l = 1; l <= N; l++)
        {
            int rightValue = B * prefix[l - 1] - C * l + C;
            m[rightValue] = m.GetValueOrDefault(rightValue, 0) + 1;
            ans += m.GetValueOrDefault(B * prefix[l] - C * l, 0);
        }

        // returning the result
        return ans;
    }

    public static void Main(string[] args)
    {
        int N = 3;
        int[] A = { 1, 2, 3 };
        int B = 1, C = 2;

        // Function call
        Console.Write(CountSubarrays(N, A, B, C));
    }
}

Javascript

function countSubarrays(N, A, B, C) {

    // Declaring the prefix vector to
      // store prefix sums of the array.
    let prefix = new Array(N + 1).fill(0);

    // Computing the prefix sums
    for (let i = 1; i <= N; i++) {
        prefix[i] = prefix[i - 1] + A[i - 1];
    }

    // ans will store the count of subarrays
    let ans = 0;
    let m = {};


    // Hashing the value on the right side
      // and then search for the value of
      // left side as key in the map
    for (let l = 1; l <= N; l++) {
        m[B * prefix[l - 1] - C * l + C] = (m[B * prefix[l - 1] - C * l + C] || 0) + 1;
        ans += m[B * prefix[l] - C * l] || 0;
    }

    // Returning the result
    return ans;
}

// Driver code
let N = 3;
let A = [1, 2, 3];
let B = 1;
let C = 2;
console.log(countSubarrays(N, A, B, C));

Output

2

Time Complexity: O(n*logn), where n is the length of the array.
Auxiliary Space: O(n).