Subarrays With Equal 0’s, 1’s and 2’s

Hard
0/120
profile
Contributed by
16 upvotes
Asked in company
Oracle

Problem statement

You are given an array containing ‘N’ integers. In the array, the elements are 0, 1 and 2. You have a simple task, find the count of non-empty subarrays containing an equal number of 0’s, 1’s and 2’s.

The subarray of ARR is a contiguous part of the array ARR, i. e. the array ARRi, ARRi+1, . . . . . , ARRj for some 0 ≤ i ≤ j ≤ N - 1.

Follow Up : 
Can you solve it in linear time and space complexities?
For Example : 
If ‘N’ = 5, ‘ARR’ = {1, 1, 0, 2, 1}

There are exactly two subarrays that contain an equal number of 0’s, 1’s and 2’s. 

The first subarray is from ‘ARR[1]’ to ‘ARR[3]’, ie: {1, 0, 2}.

The second subarray is from ‘ARR[2]’ to ‘ARR[4]’, ie: {0, 2, 1}.
Detailed explanation ( Input/output format, Notes, Images )
Input Format : 
The first line contains a single integer ‘T’ denoting the number of test cases, then each test case follows:

The first line of each test case contains a single integer ‘N’, denoting the number of elements in the array.

The second line of each test case contains N integers ‘ARR’, denoting the array elements.
Output Format : 
For each test case, print the maximum number of subarrays that contain equal number of 0’s, 1’s and 2’s.

Output for each test case will be printed in a separate line.
Note : 
You are not required to print anything; it has already been taken care of. Just implement the function.
Constraints : 
1 ≤ T ≤ 10      
1 ≤ N ≤ 1000
ARR[i] = {0, 1, 2}

Time limit: 1 sec
Sample Input 1 :
2
5
1 1 0 2 1
4
1 1 0 0
Sample Output 1 :
2
0
Explanation For Sample Input 1 :
For test case 1 :
We will print 2 because:
There are exactly two subarrays that contain an equal number of 0’s, 1’s and 2’s. 
The first subarray is from ARR[1] to ARR[3], ie: {1, 0, 2}, and the second subarray is from ARR[2] to ARR[4], ie: {0, 2, 1} 

For test case 2 : 
We will print 0 because:
The input array doesn’t contain any element equal to 2, so it’s impossible to form a non-empty subarray with an equal number of 0’s, 1’s and 2’s.
Sample Input 2 :
2
6
1 0 2 1 0 2
6
1 1 0 0 2 2
Sample Output 2 :
5
1
Hint

Try checking for all the subarrays possible in an efficient way.

Approaches (2)
Brute Force

A naive solution is to check for all the subarrays one by one, generating all the subarrays take quadratic (~O(N2)) time, and then for each generated subarray we need to iterate all the elements to calculate for the count of the number of 0’s, 1’s and 2’s. This method takes a ~O(N3) time in total.

 

We can optimize the above solution by removing the need to count the number of 0’s, 1’s and 2’s in each subarray generated, this can be done by precalculating prefix array to store the count of 0’s from ARR[0] to ARR[i] in PREF0[i], and also storing similar information for 1’s and 2’s.

 

The steps are as follows :

  1. Initialize ‘ans’ equal to 0, this variable stores the count of subarrays with an equal number of 0’s, 1’s and 2’s.
  2. Create three prefix arrays, prefix0, prefix1 and prefix2, and fill them up.
  3. Run an outer FOR loop for variable ‘i’ from 0 to N - 1, and run an inner FOR loop for variable ‘j’ from 0  to i - 1. Each time calculate the count of 0’s, 1’s and 2’s in the subarray from j + 1 to i using the prefix arrays.
    • Note that: Count of 0’s in subarray from j + 1 to i equal to prefix0[i] - prefix0[j]. Similarly, calculate the count of 1's and 2's also.
  4. Finally, return the value of ‘ans’.
Time Complexity

O( N ^ 2 ), where N denotes the size of the array.

 

We calculate prefix sums that take linear (~O(N)) time, then we generate all the subarrays using two FOR loops, there are a total of N^2 subarray generated.

Hence the time complexity is O( N^2 ).

Space Complexity

O( N ), where N denotes the size of the array.

 

Extra space is used to store the prefix sums, the space taken is of the order of the number of elements in the input array.

Hence the space complexity is O( N ).

Code Solution
(100% EXP penalty)
Subarrays With Equal 0’s, 1’s and 2’s
Full screen
Console