Intuit India | OA | Playing With 0's and 1's

Question 3 of 4


PLAYING WITH 0's and 1's


Description

Alice is playing with numbers 0's and 1's. She decided to count the number of subarrays which satisfy the following condition:

Ratio of number of 0's to number of 1's is X:Y 

She starts struggling as soon the count of numbers increase, can you help her find the solution to this problem.

Constraints

1 <= N <= 10^5
1 <= X <= 10^5
1 <= Y <= 10^S

Input Format

The first line contains N, the count of numbers. The numbers contain only 0's and l's The second line contains N space seperated integers (0 or 1)

The third line contains 2 space seperated integers X and Y.

Output Format

Print the count of subarrays which have 0 and 1 in the ratio of X:Y. Output the answer modulo (10^9 + 7).

Sample Input

4
0 1 0 1 
1 1 

Sample Output

4 

Explanation

[0 1] 0 1 
0 [1 0] 1 
0 1 [0 1]
[0 1 0 1] 

Solution :

Thanks to @zoko1124 for pointing out this solution .

#include <bits/stdc++.h>
using namespace std;
#define endl '\n'


void solve(){
int n;
cin>>n;
int a[n];
for(int i=0;i<n;i++)
    cin>>a[i];
int x,y;
cin>>x>>y;
map<int,int> mp;
mp[0]++;
int sum=0,ans=0;
for(int i=0;i<n;i++){
    if(a[i]==0)
        sum+=(-y);
    else
        sum+=x;
    if(mp.count(sum))
        ans+=mp[sum];
    mp[sum]++;
}
cout<<ans<<endl;
}
int main() {
    int t=1;
    // cin>>t;
    while(t--){
      solve();
    }
	return 0;
}
Comments (7)