Count subarrays with equal number of occurrences of two given elements in C++

We are given an array arr[] of integers. Also, two numbers A and B. The goal is to count all subarrays of arr[] such that occurrences of A and B is equal in all. If the array is [1,2,3] and A is 1 and B is 2. Subarrays will be [3], [1,2], [1,2,3].

Let us understand with examples.

Input − arr[] = { 2, 2, 1, 1, 1, 5 }; A=1, B=5

Output − Count of subarrays with equal no. of occurrences of two given elements are − 4

Explanation − Subaarays will be − [2], [2], [2,2], [1,5]. First three have 0 occurrences of 1 and 5, last one has 1 occurrence of both.

Input − arr[] = { 5,3,7,5,3 }; A=1, B=2

Output − Count of subarrays with equal no. of occurrences of two given elements are − 15

Explanation − Subaarays will be − ( that have 0 occurrences of 1 and 2 )

[5], [3], [7], [5], [3] - 5
[5,3], [3,7], [7,5], [5,3] - 4
[5,3,7], [3,7,5], [7,5,3] - 3
[5,3,7,5], [3,7,5,3] - 2
[5,3,5,7,5] - 1

Total 15 subarrays with no occurrence ( 0 occurrences) of both 1 and 2.

The approach used in the below program is as follows

We will traverse the array using two for loops to generate all subarrays possible. From i=0 to i<=size-1 and j=i to j<=size-1. Subarrays formed will be between arr[i] to arr[j]. Count frequency of A and B in each subarray. If equal then increment the count.

• Take an array arr[] of numbers.

• Function sub_EqualOccurrence(int arr[], int size, int A, int B) takes the array and returns a count of subarrays with equal no. of A and B.

• Take the initial count as 0.

• We will traverse the array using two for loops from i=0 to i<=size-1 and j=0 to j<=size-1.

• Take two variables total_A, total_B as 0 for the number of A’s and B’s in subarray arr[i] to arr[j].

• Compare arr[j] with A and B. If arr[j] is A or B then increment respective count ( total_A or total_B).

• If total_A==total_B. Increment count. ( the subarray has the same number of A’s and B’s as elements).

• At the end of both loops, return count as result.

Example

 Live Demo

#include <bits/stdc++.h>
using namespace std;
int sub_EqualOccurrence(int arr[], int size, int A, int B){
   int count = 0;
   for (int i = 0; i <= size - 1; i++){
      int total_A = 0;
      int total_B = 0;
      for (int j = i; j <= size - 1; j++){
         if (arr[j] == A){
            total_A++;
         }
         else if (arr[j] == B){
            total_B++;
         }
         if(total_A == total_B){
            count++;
         }
      }
   }
   return count;
}
// Driver code
int main(){
   int arr[] = { 2, 3, 1, 1, 4, 5 };
   int size = sizeof(arr) / sizeof(arr[0]);
   int A = 1, B = 5;
   cout<<"Count of subarrays with equal number of occurrences of two given elements are: "<<sub_EqualOccurrence(arr, size, A, B);
   return (0);
}

Output

If we run the above code it will generate the following output −

Count of subarrays with equal number of occurrences of two given elements are: 5