Count of subarrays which contains a given number exactly K times

Given an array A[] of N elements consisting of values from 1 to N with duplicates, the task is to find the total number of subarrays which contain a given number num exactly K times.

Examples:

Input: A[] = {1, 2, 1, 5, 1}, num = 1, K = 2
Output: 2
Explanation:
Subarrays {1, 2, 1, 5}, {1, 2, 1}, {2, 1, 5, 1} and {1, 5, 1} contains 1 exactly twice.

Input: A[] = {1, 5, 3, 5, 7, 5, 6, 5, 10, 5, 12, 5}, num = 5, K = 3
Output: 14

Naive Approach: A simple solution is to generate all the subarrays of the given array and count the number of subarrays in which the given number occurs exactly K times.



Time complexity: O(N2) where N is the size of the given array.

Efficient Approach:

  • Store the indices which contain the given number num.
  • Traverse the indices[] array and calculate the number of subarrays possible for every K indices.
  • The number of subarrays possible for any K indices of num is equal to the

    Product of (ith index – (i-1)th index) and ( (i + K)th index – (i+(K – 1))th index)

  • Count of all such subarrays gives the count of the total possible subarrays in the given array.

Below is the implementation of the above approach:

C++

filter_none

edit
close

play_arrow

link
brightness_4
code

// C++ program to count subarrays
// which contains a given number
// exactly K times
  
#include <bits/stdc++.h>
using namespace std;
  
// Function to return
// the count of subarrays
// which contains given
// number exactly K times
int countSubarrays(int A[], int num,
                   int K, int size)
{
    // Store the indices
    // containing num
    vector<int> indices;
    for (int i = 0; i < size; i++) {
        if (A[i] == num)
            indices.push_back(i);
    }
  
    // If the occurence of num
    // in the entire array
    // is less than K
    if (indices.size() < K)
  
        // No such subarrays are possible
        return 0;
  
    // Store the previous
    // index of num
    int prev = -1;
  
    // Store the count of
    // total subarrays
    int ans = 0;
  
    // Store the count of
    // subarrays for current
    // K occurences
    int ctr = 0;
  
    for (int i = 0;
         i <= indices.size() - K;
         i++) {
  
        ctr = indices[i] - prev;
  
        if (i < indices.size() - K) {
  
            ctr *= (indices[i + K]
                    - indices[i + K - 1]);
        }
        else {
            ctr *= ((size - 1)
                    - indices[i + K - 1] + 1);
        }
  
        ans += ctr;
        prev = indices[i];
    }
  
    return ans;
}
  
// Driver code
int main()
{
    int A[] = { 1, 5, 3, 5, 7, 5, 6,
                5, 10, 5, 12, 5 };
  
    int num = 5;
  
    int K = 3;
  
    int size = sizeof(A) / sizeof(int);
  
    cout << countSubarrays(A, num, K, size);
  
    return 0;
}

chevron_right


Java

filter_none

edit
close

play_arrow

link
brightness_4
code

// Java program to count subarrays
// which contains a given number
// exactly K times
  
import java.util.*;
public class Main {
  
    // Function to return
    // the count of subarrays
    // which contains given
    // number exactly K times
    public static int countSubarrays(
        int A[], int num,
        int K, int size)
    {
        // Store the indices
        // containing num
        ArrayList<Integer> indices
            = new ArrayList<Integer>();
  
        for (int i = 0; i < size; i++) {
            if (A[i] == num) {
                indices.add(i);
            }
        }
  
        if (indices.size() < K) {
            return 0;
        }
  
        // Store the previous
        // index of num
        int prev = -1;
  
        // Store the count of
        // total subarrays
        int ans = 0;
  
        // Store the count of
        // subarrays for current
        // K occurences
        int ctr = 0;
  
        for (int i = 0;
             i <= indices.size() - K;
             i++) {
  
            ctr = indices.get(i) - prev;
  
            if (i < indices.size() - K) {
  
                ctr *= (indices.get(i + K)
                        - indices.get(i + K - 1));
            }
            else {
                ctr *= ((size - 1)
                        - indices.get(i + K - 1) + 1);
            }
  
            ans += ctr;
            prev = indices.get(i);
        }
  
        return ans;
    }
  
    // Driver code
    public static void main(String[] args)
    {
  
        int A[] = { 1, 5, 3, 5, 7, 5,
                    6, 5, 10, 5, 12, 5 };
  
        int num = 5;
  
        int K = 3;
  
        int size = A.length;
  
        System.out.println(
            countSubarrays(A, num, K, size));
    }
}

chevron_right


Output:

14

Time Complexity: O(N), where N is the size of the array.
Space Complexity: O(N)

competitive-programming-img




My Personal Notes arrow_drop_up

Check out this Author's contributed articles.

If you like GeeksforGeeks and would like to contribute, you can also write an article using contribute.geeksforgeeks.org or mail your article to contribute@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.

Please Improve this article if you find anything incorrect by clicking on the "Improve Article" button below.