United Cows of Farmer John

USACO 2021 US Open, Gold - Problem 1.

C++

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

class Fenwick {
    private:
        vector<int> tree;
    public:
        Fenwick(int n) {tree.assign(n+1,0);}
        int query(int j) {
            int sum = 0;
            for (; j; j-=j&-j) sum += tree[j];
            return sum;
        }
        void update(int i, int d) {
            for (; i<(int)tree.size(); i+=i&-i) tree[i] += d;
        }
};

int main() {
    int n;
    cin >> n;
    vector<long> last_oc(n+1, 0);
    int b;
    Fenwick ft = Fenwick(n+2);
    long ans = 0;
    for (int r = 1; r <= n; r++) {
        cin >> b;
        if (last_oc[b] != 0) ft.update(last_oc[b], -1);
        ans += ft.query(r) - ft.query(last_oc[b]);
        last_oc[b] = r;
        ft.update(r,1);
    }
    cout << ans << endl;
}