h9IQVp - Online C++0x Compiler & Debugging Tool

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

struct Fenwick {
    int n;
    vector<int> bit;
    Fenwick(int n=0){init(n);}
    void init(int n_){
        n = n_;
        bit.assign(n+1, 0);
    }
    void add(int i, int v){
        for(; i<=n; i += i&-i) bit[i] += v;
    }
    int sumPrefix(int i){
        int s=0;
        for(; i>0; i -= i&-i) s += bit[i];
        return s;
    }
    int sumRange(int l, int r){
        if(r<l) return 0;
        return sumPrefix(r) - sumPrefix(l-1);
    }
};

int main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int N;
    if(!(cin >> N)) return 0;
    vector<int> A(N+1);
    int maxA = 0;
    for(int i=1;i<=N;i++){ cin >> A[i]; maxA = max(maxA, A[i]); }

    // positions per value
    int V = maxA;
    vector<vector<int>> pos(V+1);
    for(int i=1;i<=N;i++){
        pos[A[i]].push_back(i);
    }

    const int B = 700; // threshold ~ sqrt(N)
    ll total = (ll)N * (N+1) / 2;
    ll bad = 0; // number of subarrays that have a majority (to be subtracted)

    // 1) Heavy values (freq > B)
    for(int val = 1; val <= V; ++val){
        int k = (int)pos[val].size();
        if(k <= B) continue;
        // build prefix S: +1 if A[i]==val else -1
        vector<int> S(N+1);
        S[0] = 0;
        for(int i=1;i<=N;i++){
            S[i] = S[i-1] + (A[i]==val ? 1 : -1);
        }
        // coordinate compress S values
        vector<int> comp = S;
        sort(comp.begin(), comp.end());
        comp.erase(unique(comp.begin(), comp.end()), comp.end());
        int m = (int)comp.size();
        Fenwick fw(m+5);
        // count pairs i<j with S[j] > S[i]
        for(int i=0;i<=N;i++){
            int id = int(lower_bound(comp.begin(), comp.end(), S[i]) - comp.begin()) + 1;
            // number of previous S < S[i] = sum of indices < id
            bad += fw.sumPrefix(id-1);
            fw.add(id, 1);
        }
    }

    // 2) Light values (freq <= B)
    // For each value, consider windows of t consecutive occurrences
    for(int val = 1; val <= V; ++val){
        int k = (int)pos[val].size();
        if(k == 0 || k > B) continue;
        // create augmented pos with boundaries
        vector<int> p;
        p.reserve(k+2);
        p.push_back(0);
        for(int x: pos[val]) p.push_back(x);
        p.push_back(N+1);
        // p indices: 0..k+1, actual occurrences at 1..k
        // iterate t = number of occurrences inside subarray
        for(int t = 1; t <= k; ++t){
            // slide window of length t over occurrences
            for(int i = 1; i + t - 1 <= k; ++i){
                int left_prev = p[i-1];
                int left_first = p[i];
                int right_last = p[i+t-1];
                int right_next = p[i+t];
                // L in [left_prev+1, left_first]
                int Lstart = left_prev + 1;
                int Lend = left_first;
                if(Lstart > Lend) continue;
                // R min = right_last
                int Rmin = right_last;
                int Rmax_cap = right_next - 1;
                if(Rmin > Rmax_cap) continue; // no valid R at all
                // for given L, Rmax = min(Rmax_cap, L + 2t - 2)
                // Count_R = max(0, Rmax - Rmin + 1)
                // We sum Count_R for L in [Lstart..Lend]
                long long Cconst = (long long)(2*t - 1 - Rmin); // Count_R = L + Cconst when L+2t-2 <= Rmax_cap
                // D = Rmax_cap - (2t-2) : threshold L where L+2t-2 >= Rmax_cap
                int D = Rmax_cap - (2*t - 2);
                // range1: L in [L1..L2] where L <= min(Lend, D) and also L >= Lstart and L >= L_lower_bound so that Count_R>0
                int L1 = max(Lstart, Rmin - (2*t - 2)); // ensure L+2t-2 >= Rmin -> Count_R positive
                int L2 = min(Lend, D);
                long long sum = 0;
                if(L1 <= L2){
                    // sum_{L=L1..L2} (L + Cconst) = sum L + (L2-L1+1)*Cconst
                    long long cnt = L2 - L1 + 1;
                    long long sumL = (long long)(L1 + L2) * cnt / 2;
                    sum += sumL + cnt * Cconst;
                }
                // range2: L in [max(Lstart, D+1) .. Lend], Count_R = Rmax_cap - Rmin +1 (if positive)
                int L3 = max(Lstart, D+1);
                int L4 = Lend;
                if(L3 <= L4){
                    long long cnt = L4 - L3 + 1;
                    long long valR = Rmax_cap - Rmin + 1;
                    if(valR > 0) sum += cnt * valR;
                }
                if(sum > 0) bad += sum;
            }
        }
    }

    ll good = total - bad;
    cout << good << '\n';
    return 0;
}