Count of subarrays with average K

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 <bits/stdc++.h>
using namespace std;

// Function to count subarray having average
// exactly equal to K
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 code
int main()
{
    // Given Input
    int K = 6;
    int arr[] = { 12, 5, 3, 10, 4, 8, 10, 12, -6, -1 };
    int N = sizeof(arr) / sizeof(arr[0]);

    // Function Call
    cout << countKAverageSubarrays(arr, N, K);
}

Java

// Java implementation of above approach
import java.io.*;

class GFG{

// Function to count subarray having average
// exactly equal to K
static 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 code
public 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 K
def 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 code
if __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 approach
using System;
using System.Collections.Generic;

class GFG{

// Function to count subarray having average
// exactly equal to K
static 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 code
public 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

<script>
// Javascript implementation of above approach

// Function to count subarray having average
// exactly equal to K
function countKAverageSubarrays(arr, n, k)
{

  // To Store the final answer
  let res = 0;

  // Calculate all subarrays
  for (let L = 0; L < n; L++)
  {
    let sum = 0;
    for (let R = L; R < n; R++)
    {

      // Calculate required average
      sum += arr[R];
      let len = R - L + 1;

      // Check if average
      // is equal to k
      if (sum % len == 0)
      {
        let avg = sum / len;

        // Required average found
        if (avg == k)
          // Increment res
          res++;
      }
    }
  }
  return res;
}

// Driver code

// Given Input
let K = 6;
let arr = [12, 5, 3, 10, 4, 8, 10, 12, -6, -1];
let N = arr.length;

// Function Call
document.write(countKAverageSubarrays(arr, N, K));

// This code is contributed by saurabh_jaiswal.
</script>

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 <bits/stdc++.h>
using namespace std;

// Function to count subarray having average
// exactly equal to K
int countKAverageSubarrays(int arr[], int n, int k)
{
    int result = 0, curSum = 0;

    // Store the frequency of prefix
    // sum of the array arr[]
    unordered_map<int, int> 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 code
int main()
{
    // Given Input
    int K = 6;
    int arr[] = { 12, 5, 3, 10, 4, 8, 10, 12, -6, -1 };
    int N = sizeof(arr) / sizeof(arr[0]);

    // Function Call
    cout << countKAverageSubarrays(arr, N, K);
}

Java

// Java program for the above approach
import java.util.*;

class GFG{

// Function to count subarray having average
// exactly equal to K
static int countKAverageSubarrays(int[] arr, int n,
                                  int k)
{
    int result = 1, curSum = 0;

    // Store the frequency of prefix
    // sum of the array arr[]
    HashMap<Integer,
            Integer> mp = new HashMap<Integer,
                                      Integer>();

    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
          if(mp.containsKey(curSum))
            mp.put(curSum, mp.get(curSum) + 1);
        else
            mp.put(curSum, 1);
    }

    // Return result
    return result;
}

// Driver Code
public 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 K
def 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 code
if __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 approach
using System;
using System.Collections.Generic;


public class GFG{

// Function to count subarray having average
// exactly equal to K
static int countKAverageSubarrays(int[] arr, int n,
                                  int k)
{
    int result = 1, curSum = 0;

    // Store the frequency of prefix
    // sum of the array []arr
    Dictionary<int,
            int> mp = new Dictionary<int,
                                      int>();

    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 Code
public 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

<script>

        // JavaScript Program for the above approach

        // Function to count subarray having average
        // exactly equal to K
        function countKAverageSubarrays(arr, n, k) {
            let result = 0, curSum = 0;

            // Store the frequency of prefix
            // sum of the array arr[]
            let mp = new Map();

            for (let 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.has(curSum)) {
                    result += mp.get(curSum);
                }


                // Increment the frequency
                // of curSum
                if (mp.has(curSum)) {
                    mp.set(curSum, mp.get(curSum) + 1);
                }
                else {
                    mp.set(curSum, 1);
                }

            }

            // Return result
            return result;
        }

        // Driver code

        // Given Input
        let K = 6;
        let arr = [12, 5, 3, 10, 4, 8, 10, 12, -6, -1];
        let N = arr.length;

        // Function Call
        document.write(countKAverageSubarrays(arr, N, K));

    // This code is contributed by Potta Lokesh

</script>

Output

4

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