Count of subarrays with average K

• Difficulty Level : Medium
• Last Updated : 11 Aug, 2021

Given an array arr[] of size N, the task is to count the number of subarrays having an average exactly equal to k.

Examples:

Input: arr[ ] = {1, 4, 2, 6, 10}, N = 6, K = 4
Output: 3
Explanation: The subarrays with an average equal to 4 are {4}, {2, 6}, {4, 2, 6}.

Input: arr[ ] = {12, 5, 3, 10, 4, 8, 10, 12, -6, -1}, N = 10, K = 6
Output: 4

Naive Approach: The simplest approach to solve the problem is to traverse all the subarrays and calculate their average. If their average is K, then increase the answer.

Below is the implementation of the naive approach:

C++

 // C++ implementation of above approach#include using namespace std; // Function to count subarray having average// exactly equal to Kint countKAverageSubarrays(int arr[], int n, int k){     // To Store the final answer    int res = 0;     // Calculate all subarrays    for (int L = 0; L < n; L++) {        int sum = 0;        for (int R = L; R < n; R++) {            // Calculate required average            sum += arr[R];            int len = (R - L + 1);             // Check if average            // is equal to k            if (sum % len == 0) {                int avg = sum / len;                 // Required average found                if (avg == k)                     // Increment res                    res++;            }        }    }    return res;} // Driver codeint main(){    // Given Input    int K = 6;    int arr[] = { 12, 5, 3, 10, 4, 8, 10, 12, -6, -1 };    int N = sizeof(arr) / sizeof(arr);     // Function Call    cout << countKAverageSubarrays(arr, N, K);}

Java

 // Java implementation of above approachimport java.io.*; class GFG{     // Function to count subarray having average// exactly equal to Kstatic int countKAverageSubarrays(int arr[], int n,                                  int k){         // To Store the final answer    int res = 0;     // Calculate all subarrays    for(int L = 0; L < n; L++)    {        int sum = 0;        for(int R = L; R < n; R++)        {                         // Calculate required average            sum += arr[R];            int len = (R - L + 1);             // Check if average            // is equal to k            if (sum % len == 0)            {                int avg = sum / len;                 // Required average found                if (avg == k)                     // Increment res                    res++;            }        }    }    return res;} // Driver codepublic static void main(String[] args){         // Given Input    int K = 6;    int arr[] = { 12, 5, 3, 10, 4,                   8, 10, 12, -6, -1 };    int N = arr.length;     // Function Call    System.out.print(countKAverageSubarrays(arr, N, K));}} // This code is contributed by shivanisinghss2110

Python3

 # Python 3 implementation of above approach # Function to count subarray having average# exactly equal to Kdef countKAverageSubarrays(arr, n, k):    # To Store the final answer    res = 0     # Calculate all subarrays    for L in range(n):        sum = 0        for R in range(L,n,1):            # Calculate required average            sum += arr[R]            len1 = (R - L + 1)             # Check if average            # is equal to k            if (sum % len1 == 0):                avg = sum // len1                 # Required average found                if (avg == k):                     # Increment res                    res += 1                 return res # Driver codeif __name__ == '__main__':    # Given Input    K = 6    arr = [12, 5, 3, 10, 4, 8, 10, 12, -6, -1]    N = len(arr)     # Function Call    print(countKAverageSubarrays(arr, N, K))     # This code is contributed by SURENDRA_GANGWAR.

C#

 // C# implementation of above approachusing System;using System.Collections.Generic; class GFG{ // Function to count subarray having average// exactly equal to Kstatic int countKAverageSubarrays(int []arr, int n, int k){     // To Store the final answer    int res = 0;     // Calculate all subarrays    for (int L = 0; L < n; L++)    {        int sum = 0;        for (int R = L; R < n; R++)        {                       // Calculate required average            sum += arr[R];            int len = (R - L + 1);             // Check if average            // is equal to k            if (sum % len == 0) {                int avg = sum / len;                 // Required average found                if (avg == k)                     // Increment res                    res++;            }        }    }    return res;} // Driver codepublic static void Main(){    // Given Input    int K = 6;    int []arr = { 12, 5, 3, 10, 4, 8, 10, 12, -6, -1 };    int N = arr.Length;     // Function Call    Console.Write(countKAverageSubarrays(arr, N, K));}} // This code is contributed by bgangwar59.

Javascript


Output
4

Time Complexity: O(N^2)
Auxiliary Space: O(1)

Efficient Approach: An efficient solution is based on the observations below:

Let there be a subarray [L, R] whose average is equal to K, then
=> K = average[L, R] = sum[0, R] – sum[0, L-1] / (R – L + 1)
=> (R – L + 1) * K = sum[0, R] – sum[0, L – 1]
=> R * k – (L – 1)* K = sum[0, R] – sum[0, L – 1]
=> sum[0, R] – R * k = sum[0, L – 1]  – (L – 1)* K

If every element is decreased by K, then the average will also decrease by K. Therefore, the average can be reduced to zero, so the problem becomes finding the number of subarrays having average equals zero.
The average zero is possible only if:
sum[0, R] – sum[0, L-1] / (R – L + 1) = 0
=> sum[0, R] = sum[0, L-1]

Follow the steps below to solve this problem :

• Initialize a map say, mp to store the frequency of prefix sum of the array arr[].
• Initialize a variable, say, curSum and result as 0.
• Iterate in the range[0, N-1] using the variable i:
• Subtract K from current element, then add it to curSum.
• If curSum is 0, subarray having an average equal to 0 is found, so increment the result by 1.
• If curSum has previously occurred before using a map. If it has occurred before then add the number of times it has occurred before to the result, then increase the frequency of curSum using the map.
• After completing the above steps, print the result as the answer.

Below is the implementation of the above approach:

C++

 // C++ program for the above approach#include using namespace std; // Function to count subarray having average// exactly equal to Kint countKAverageSubarrays(int arr[], int n, int k){    int result = 0, curSum = 0;     // Store the frequency of prefix    // sum of the array arr[]    unordered_map mp;     for (int i = 0; i < n; i++) {        // Subtract k from each element,        // then add it to curSum        curSum += (arr[i] - k);         // If curSum is 0 that means        // sum[0...i] is 0 so increment        // res        if (curSum == 0)            result++;         // Check if curSum has occurred        // before and if it has occurred        // before, add it's frequency to        // res        if (mp.find(curSum) != mp.end())            result += mp[curSum];         // Increment the frequency        // of curSum        mp[curSum]++;    }     // Return result    return result;} // Driver codeint main(){    // Given Input    int K = 6;    int arr[] = { 12, 5, 3, 10, 4, 8, 10, 12, -6, -1 };    int N = sizeof(arr) / sizeof(arr);     // Function Call    cout << countKAverageSubarrays(arr, N, K);}

Java

 // Java program for the above approachimport java.util.*; class GFG{     // Function to count subarray having average// exactly equal to Kstatic int countKAverageSubarrays(int[] arr, int n,                                  int k){    int result = 1, curSum = 0;     // Store the frequency of prefix    // sum of the array arr[]    HashMap mp = new HashMap();     for(int i = 0; i < n; i++)    {                 // Subtract k from each element,        // then add it to curSum        curSum += (arr[i] - k);         // If curSum is 0 that means        // sum[0...i] is 0 so increment        // res        if (curSum == 0)            result++;         // Check if curSum has occurred        // before and if it has occurred        // before, add it's frequency to        // res        if (mp.containsKey(curSum))            result += mp.get(curSum);         // Increment the frequency        // of curSum        mp.put(curSum, 1);    }     // Return result    return result;} // Driver Codepublic static void main(String[] args){         // Given Input    int K = 6;    int[] arr = { 12, 5, 3, 10, 4,                  8, 10, 12, -6, -1 };    int N = arr.length;     // Function Call    System.out.print(countKAverageSubarrays(arr, N, K));}} // This code is contributed by sanjoy_62

Python3

 # Python Program for the above approach # Function to count subarray having average# exactly equal to Kdef countKAverageSubarrays(arr, n, k):       result = 0    curSum = 0         # Store the frequency of prefix    # sum of the array arr[]    mp = dict()         for i in range(0, n):               # Subtract k from each element,        # then add it to curSum        curSum += (arr[i] - k)                 # If curSum is 0 that means        # sum[0...i] is 0 so increment        # res        if (curSum == 0):            result += 1                     # Check if curSum has occurred        # before and if it has occurred        # before, add it's frequency to        # res        if curSum in mp:            result += mp[curSum]                     # Increment the frequency        # of curSum        if curSum in mp:            mp[curSum] += 1        else:            mp[curSum] = 1                 # Return result    return result  # Driver codeif __name__ == '__main__':       # Given Input    K = 6    arr = [12, 5, 3, 10, 4, 8, 10, 12, -6, -1]    N = len(arr)     # Function Call    print(countKAverageSubarrays(arr, N, K)) # This code is contributed by MuskanKalra1

C#

 // C# program for the above approachusing System;using System.Collections.Generic;  public class GFG{     // Function to count subarray having average// exactly equal to Kstatic int countKAverageSubarrays(int[] arr, int n,                                  int k){    int result = 1, curSum = 0;     // Store the frequency of prefix    // sum of the array []arr    Dictionary mp = new Dictionary();     for(int i = 0; i < n; i++)    {                 // Subtract k from each element,        // then add it to curSum        curSum += (arr[i] - k);         // If curSum is 0 that means        // sum[0...i] is 0 so increment        // res        if (curSum == 0)            result++;         // Check if curSum has occurred        // before and if it has occurred        // before, add it's frequency to        // res        if (mp.ContainsKey(curSum))            result += mp[curSum];        else        // Increment the frequency        // of curSum        mp.Add(curSum, 1);    }     // Return result    return result;} // Driver Codepublic static void Main(String[] args){         // Given Input    int K = 6;    int[] arr = { 12, 5, 3, 10, 4,                  8, 10, 12, -6, -1 };    int N = arr.Length;     // Function Call    Console.Write(countKAverageSubarrays(arr, N, K));}}  // This code contributed by shikhasingrajput

Javascript


Output
4

Time Complexity: O(N)
Auxiliary Space: O(N)

My Personal Notes arrow_drop_up