Count of Subarray with B times Sum equal to C times Length
Last Updated :
19 Sep, 2023
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++
#include <bits/stdc++.h>
using namespace std;
int countSubarrays( int N, int A[], int B, int C)
{
vector< long long > prefix(N + 1);
for ( int i = 1; i <= N; i++) {
prefix[i] = prefix[i - 1] + A[i - 1];
}
long long ans = 0;
map< long long , int > m;
for ( int l = 1; l <= N; l++) {
m[B * prefix[l - 1] - C * l + C] += 1;
ans += m[B * prefix[l] - C * l];
}
return ans;
}
int main()
{
int N = 3;
int A[] = { 1, 2, 3 };
int B = 1, C = 2;
cout << countSubarrays(N, A, B, C);
return 0;
}
|
Java
import java.io.*;
import java.util.*;
class GFG {
static int countSubarrays( int N, int [] A, int B, int C)
{
int [] prefix = new int [N + 1 ];
for ( int i = 1 ; i <= N; i++) {
prefix[i] = prefix[i - 1 ] + A[i - 1 ];
}
int ans = 0 ;
Map<Integer, Integer> m = new HashMap<>();
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 );
}
return ans;
}
public static void main(String[] args)
{
int N = 3 ;
int [] A = { 1 , 2 , 3 };
int B = 1 , C = 2 ;
System.out.print(countSubarrays(N, A, B, C));
}
}
|
Python3
def countSubarrays(N, A, B, C):
prefix = [ 0 ] * (N + 1 )
for i in range ( 1 , N + 1 ):
prefix[i] = prefix[i - 1 ] + A[i - 1 ]
ans = 0
m = {}
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 )
return ans
N = 3
A = [ 1 , 2 , 3 ]
B = 1
C = 2
print (countSubarrays(N, A, B, C))
|
C#
using System;
using System.Collections.Generic;
public class GFG
{
static int CountSubarrays( int N, int [] A, int B, int C)
{
int [] prefix = new int [N + 1];
for ( int i = 1; i <= N; i++)
{
prefix[i] = prefix[i - 1] + A[i - 1];
}
int ans = 0;
Dictionary< int , int > m = new Dictionary< int , int >();
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);
}
return ans;
}
public static void Main( string [] args)
{
int N = 3;
int [] A = { 1, 2, 3 };
int B = 1, C = 2;
Console.Write(CountSubarrays(N, A, B, C));
}
}
|
Javascript
function countSubarrays(N, A, B, C) {
let prefix = new Array(N + 1).fill(0);
for (let i = 1; i <= N; i++) {
prefix[i] = prefix[i - 1] + A[i - 1];
}
let ans = 0;
let m = {};
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;
}
return ans;
}
let N = 3;
let A = [1, 2, 3];
let B = 1;
let C = 2;
console.log(countSubarrays(N, A, B, C));
|
Time Complexity: O(n*logn), where n is the length of the array.
Auxiliary Space: O(n).
Share your thoughts in the comments
Please Login to comment...