Code : Count Winning Subarrays

#include <bits/stdc++.h>
using namespace std;

void solve(vector<int> &a) {
    int n = a.size();

    long long ans = 0;
    int lst = -1;
    // How many winning subarrays end at index i?
    for (int i = 0; i < n; i++) {
        if (i > 0 && a[i] == 1 && a[i - 1] == 1) {
            lst = max(lst, i - 1);
        }
        if (i > 1 && a[i] == 1 && a[i - 1] == 0 && a[i - 2] == 1) {
            lst = max(lst, i - 2);
        }

        // |[0, lst]| are the safe indices.
        ans += lst + 1;
        ans += (a[i] == 1);
    }
    cout << ans << "\n";
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int t;
    cin >> t;
    while (t--) {
        int n;
        cin >> n;
        vector<int> a(n);
        for (int i = 0; i < n; i++) {
            cin >> a[i];
        }
        solve(a);
    }
}