find the total number of subarrays with the ratio of 0's and 1's equal to x:y

Question

question

given an array of elements 0, 1, 2 with find the total number of subarrays with the ratio of 0's and 1's equal to x:y.

input

---

5

---

1 1

---

0 1 2 0 1

---

output 6

---

\\5 is the size of array 0 1 2 0 1 are elements of the array 1 1 is x and y and now we have to find the subarrays whose counts of 0's and 1's ratio is equal to x and y that is 1 1 \\

---

here is my approach but it is not giving correct and it gives output 7 rather than 6

---

#include<bits/stdc++.h>

using namespace std;

int n, x, y;
vector<int> a;
vector<long long> prefix;
map<long long, int> freq;

int main() {
    cin >> n;
    cin >> x >> y;

    a.resize(n);
    for (int i = 0; i < n; i++) {
        cin >> a[i];
        if (a[i]==0) a[i] = y;
        else if( a[i]==0){
            a[i]=0;
        }
        else a[i] = -x;
    }

    prefix.resize(n+1); prefix[0] = 0;
    for (int i = 0; i < n; i++) {
        prefix[i+1] = prefix[i] + a[i];
    }

    for (int i = 0; i < n+1; i++) freq[prefix[i]]++;
    long long ans = 0;
    for (pair<long long, int> p : freq) {
        ans += (long long) p.second * (p.second-1) / 2;
    }
    cout << ans << endl;
}

Answer 1

First, the order you read your inputs is different than what you describe, then your transform of the input values make no sense.

//... 

cin >> n;
cin >> x >> y;  // program expects x then y, then contents of array.

a.resize(n);
for (int i = 0; i < n; i++) {
    cin >> a[i];
    if (a[i]==0) a[i] = y;
    else if( a[i]==0){       ∕/ case a[i] == 0 was already tested for
                             // on line above. 
        a[i]=0;
    }
    else a[i] = -x;
}

//...  

You should sort this out first.