anupomkar's blog

By anupomkar, history, 19 months ago, In English

Given a binary array, find the total number of subarrays with the ratio of 0's and 1's equal to x:y.


n = arr.length



Example: arr = [0,1,0,1] x = 1 y = 1

Output: 4

Explanation: The four possible subarrays are as follows:

[0 1] 0 1 --> [0,1]

0 [1 0] 1 --> [1,0]

0 1 [0 1] --> [0,1]

[0 1 0 1] --> [0,1,0,1]

With the given constraints, expecting solution no worse than O(n)

P.s: Maybe the problem can be solved using the sliding window technique. But unable to apply it due to ratio constraint.

  • Vote: I like it
  • +8
  • Vote: I do not like it

| Write comment?
19 months ago, # |
Rev. 2   Vote: I like it 0 Vote: I do not like it

Upd 1: Didn't notice O(n) constraint. This solution is O(n^2). It can probably be optimized to O(n*log n) if you first calculate the answer for smallest window size and use this result to calculate for large sizes.

Pretty sure that you got the right idea: sliding window works. For ratio constraints: just generate all possible windows: they will have sizes of (a*x+a*y), obviously (a*x+a*y)<=n. So for n == 4 and x = y = 1 you'll get windows of sizes:

1*1+1*1 = 2
2*1+2*1 = 4

(3*1+3*1) = 6 won't work as 6 > 4(=n).

For n = 9 and 2:3 you'll get windows of size 5.

Notice that you'll have to divide x and y by their GCD before applying this. So, for example, if you get x=4, y=8, you should set x=1 (4/gcd(4,8)) and y=2 (8/gcd(4,8)).

19 months ago, # |
Rev. 3   Vote: I like it +4 Vote: I do not like it

First, think about how to solve the problem for x = y = 1. One possible solution is this: replace 0 with -1, then calculate the prefix sum for the array (including an extra 0 for the beginning). For example, if the original array is $$$[0, 1, 0, 1]$$$, we replace 0 with -1 to get $$$[-1, 1, -1, 1]$$$, then calculate prefix sums to get $$$[0, -1, 0, -1, 0]$$$. Notice that if two indices have the same prefix sum, then the range between them must have had the same number of 0's and 1's. So we can count how many times each prefix value appears (call it cnt) then add $$$\frac{cnt \cdot (cnt - 1)}{2}$$$ to the answer because that's how many ways there are to pick two of those endpoints.

In order to generalize this to ratio x:y, we can do the following: replace 0 with y, and 1 with -x, then do the same thing as before. We can prove that the sum of a range being 0 is enough to guarantee that the ratio of 0's to 1's is x:y. Let's say there are $$$a$$$ 0's and $$$b$$$ 1's. Then the sum on the range is $$$a(y) + b(-x) = 0$$$. From there, we can get $$$ay = bx$$$, and $$$\frac{a}{b} = \frac{x}{y}$$$, which is exactly the original ratio condition we needed to satisfy.

My code has time complexity $$$O(n \cdot log(n))$$$, which should pass. It can be optimized to $$$O(n)$$$ with unordered_map, but that's unnecessary here.