Last Updated: 15 Mar, 2021

Count Subarrays

Easy

Asked in companies

Solve now

Problem statement

You are given an array/list consisting of 0 and 1 only. Your task is to find the sum of the number of subarrays that contains only 1 and the number of subarrays that contains only 0.

An array 'C' is a subarray of array 'D' if 'C' can be obtained from 'D' by deletion of several elements from the beginning and several elements from the end. Example :

Let 'ARR' = [1,0,0] then the possible subarrays of 'ARR' will be: {1}, {0}, {0}, {1,0}, {0,0}, {1,0,0}.

Example

If the given array 'ARR' = [1,0,0,0,1,1,0,1]
Then the number of 1’s subarrays will be 5. [{1},{1},{1},{1,1},{1}]
And the number of 0’s subarrays will be 7. [{0},{0},{0},{0,0},{0,0},{0,0,0},{0}]
So our answer will be 5 + 7 = 12.

Input Format:

The first line of input contains an integer ‘T’ denoting the number of test cases to run. Then the test case follows.

The first line of each test case contains an integer  ‘N’ representing the array’s size.

The second line of each test case contains N space-separated integers representing the array’s elements.

Output Format:

For each test case print, the sum of the number of 1’s subarrays and the number of 0’s subarrays.

For each test case print output in a separate line.

Note:

You don’t have to take any input or print anything; it already has been taken care of. Just implement the function.

Constraints:-

1 <= T <= 100
1 <= N <= 5000
0 <= ARR[i] <= 1

Where ARR[i] denotes the elements of the array.

Time Limit: 1 sec

Solve now

Approaches

01 Approach

Approach: We will take a variable to ‘COUNT1’ = 0 to store the count of 1, and another variable ‘SUM1’ = 0 to store the total number of subarrays having 1 only. Then we will iterate on each element of the array. If it is 1 then we will increase count1 by 1, else we will add (COUNT1)*(COUNT1 + 1)/2 to the ‘SUM1' i.e count of subarrays, and set ‘COUNT1’ = 0 again.

We will do the same for 0 also, by taking variables ‘COUNT0’ = 0, and ‘SUM0’ = 0.

After that, we will return ‘SUM1’ + ‘SUM0’.

Algorithm:

Initalise variable ‘COUNT1’ = 0 and ‘SUM1’ = 0.
Now for each ‘i’ from 0 to ‘N - 1':
if ‘ARR[i]’ equals 1
Increment ‘COUNT1’ by 1;E
Else
‘SUM1’ += ‘COUNT1’ * ('COUNT1' + 1) /2;
‘COUNT1’ = 0.
Now for the last set of subarrays we add ‘COUNT1 * (COUNT1 + 1) / 2’ to the ‘SUM1’.
Repeat the same for finding the count of subarrays for 0, by taking the variable let’s say ‘COUNT1’ = 0 and ‘SUM0’ = 0.
We finally return ‘SUM1’ + ‘SUM0’.

Solve now