Submission #1159621

#TimeUsernameProblemLanguageResultExecution timeMemory
1159621panIzbori (COCI22_izbori)C++17
0 / 110
77 ms14264 KiB
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

// BIT implementation for range updates & range queries:
void bit_update(vector<ll>& bit, int idx, ll val) {
    for(; idx < (int)bit.size(); idx += idx & -idx)
        bit[idx] += val;
}
ll bit_query(const vector<ll>& bit, int idx) {
    ll res = 0;
    for(; idx > 0; idx -= idx & -idx)
        res += bit[idx];
    return res;
}

// These functions implement the standard two‐BIT method.
// For an array A, to support range updates of A and range sum queries,
// we maintain two BITs (bit1 and bit2) and use the formula:
//   prefix_sum(x) = bit_query(bit1, x) * x - bit_query(bit2, x)
// and range sum [l, r] = prefix_sum(r) - prefix_sum(l-1).
void update_range(vector<ll>& bit1, vector<ll>& bit2, int l, int r, ll v) {
    bit_update(bit1, l, v);
    bit_update(bit1, r + 1, -v);
    bit_update(bit2, l, v * (l - 1));
    bit_update(bit2, r + 1, -v * r);
}
ll query_prefix(const vector<ll>& bit1, const vector<ll>& bit2, int idx) {
    return bit_query(bit1, idx) * idx - bit_query(bit2, idx);
}
ll query_range(const vector<ll>& bit1, const vector<ll>& bit2, int l, int r) {
    return query_prefix(bit1, bit2, r) - query_prefix(bit1, bit2, l - 1);
}

struct Interval {
    int l, r;
};

int main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    
    int n;
    cin >> n;
    vector<ll> arr(n);
    vector<ll> dis;
    // dish[i] will hold all indices where the i-th distinct value occurs.
    vector<vector<ll>> dish;
    
    for (int i = 0; i < n; i++){
        cin >> arr[i];
        dis.push_back(arr[i]);
    }
    sort(dis.begin(), dis.end());
    dis.erase(unique(dis.begin(), dis.end()), dis.end());
    
    dish.resize(dis.size());
    // Map each distinct value to its bucket.
    unordered_map<ll,int> mp;
    for (int i = 0; i < (int)dis.size(); i++){
        mp[dis[i]] = i;
    }
    for (int i = 0; i < n; i++){
        int idx = mp[arr[i]];
        dish[idx].push_back(i);
    }
    
    // BIT coordinate range: we need to cover indices in [-n-5, n+5].
    // We shift coordinates by OFFSET = n+5 so they become positive.
    int OFFSET = n + 5;
    int M = 2 * n + 11; // total number of coordinates.
    // Our BITs are 1-indexed and of size M+1.
    vector<ll> bitC1(M + 1, 0), bitC2(M + 1, 0); // for counts (the "counter")
    vector<ll> bitW1(M + 1, 0), bitW2(M + 1, 0); // for weighted sums
    
    // We'll record all update intervals so we can later revert them.
    vector<Interval> logs;
    ll ans = 0;
    
    // Process each distinct value.
    for (size_t i = 0; i < dis.size(); i++){
        logs.clear();
        int curOffset = 0;
        vector<ll>& now = dish[i];
        // Ensure the indices are in order.
        sort(now.begin(), now.end());
        for (size_t j = 0; j < now.size(); j++){
            int start = (j == 0 ? 0 : now[j - 1] + 1);
            int len = now[j] - start;
            if(len >= 1) {
                // Map the range [curOffset, curOffset+len-1] to BIT indices.
                int L = curOffset + OFFSET + 1;
                int R = curOffset + len - 1 + OFFSET + 1;
                update_range(bitC1, bitC2, L, R, 1);
                update_range(bitW1, bitW2, L, R, 1);
                logs.push_back({L, R});
            }
            curOffset += len;
            int pos = curOffset + OFFSET + 1;
            update_range(bitC1, bitC2, pos, pos, 1);
            update_range(bitW1, bitW2, pos, pos, 1);
            logs.push_back({pos, pos});
            curOffset--;
            int end = (j + 1 < now.size() ? now[j + 1] - 1 : n - 1);
            len = end - now[j] + 1;
            if(len > 0) {
                int Lq = curOffset + 1 + OFFSET + 1;
                int Rq = curOffset + len + OFFSET + 1;
                ll cnt = query_range(bitC1, bitC2, Lq, Rq);
                ll wsum = query_range(bitW1, bitW2, Lq, Rq);
                ans += wsum - (ll)curOffset * cnt;
                int Lq2 = curOffset + len + 1 + OFFSET + 1;
                int Rq2 = n + 1 + OFFSET + 1; // original code queries up to n+1.
                if(Lq2 <= Rq2) {
                    ll cnt2 = query_range(bitC1, bitC2, Lq2, Rq2);
                    ans += cnt2 * len;
                }
            }
        }
        // Revert the updates made in this iteration.
        for(auto &intr : logs) {
            update_range(bitC1, bitC2, intr.l, intr.r, -1);
            update_range(bitW1, bitW2, intr.l, intr.r, -1);
        }
    }
    
    cout << ans << "\n";
    return 0;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...