Submission #1161692

#	Time	Username	Problem	Language	Result	Execution time	Memory
1161692		antonn	Diversity (CEOI21_diversity)	C++20	100 / 100	6750 ms	3556 KiB

diversity

#include <bits/stdc++.h>

#define F first
#define S second
 
using namespace std;
using ll = long long;
using pi = pair<int, int>;
using vi = vector<int>;
 
template<class T> bool ckmin(T& a, T b) { return b < a ? a = b, true : false; }
template<class T> bool ckmax(T& a, T b) { return a < b ? a = b, true : false; }

const int N = 3e5 + 7;
const int K = 550;

int a[N], cnt[N];

int f[N], small[K], taken[K];
multiset<int> s;

void change(int x, int sgn) {
    if (f[x] >= K) {
        s.erase(s.find(f[x]));
    } else {
        --small[f[x]];
    }
    f[x] += sgn;
    if (f[x] >= K) {
        s.insert(f[x]);
    } else {
        ++small[f[x]];
    }
}

ll gauss(ll n) {
    return n * (n + 1) / 2;
}

ll sos(ll n) {
    return n * (n + 1) * (2 * n + 1) / 6;
}

ll get(ll sum, ll x, ll total, ll cnt) {
    ll ans = 0;
    ans = total * (sum * cnt + x * gauss(cnt - 1));
    ans -= sum * (sum * cnt + x * 2 * gauss(cnt - 1));
    ans -=  x * x * sos(cnt - 1);
    return ans;
}

int main() {
    ios::sync_with_stdio(0); cin.tie(0);
    
    int n, q; cin >> n >> q;
    for (int i = 1; i <= n; ++i) cin >> a[i];
    
    vector<array<int, 3>> qs;
    vector<ll> sol(q + 1);
    for (int i = 1; i <= q; ++i) {
        int l, r; cin >> l >> r;
        sol[i] += (ll) (r - l + 1) * (r - l + 2) / 2;
        qs.push_back({l, r, i});
    }
    sort(qs.begin(), qs.end(), [&](array<int, 3> a, array<int, 3> b) { return a[0] / K < b[0] / K || (a[0] / K == b[0] / K && a[1] < b[1]); });
    
    int l = 1, r = 0;
    for (int i = 0; i < q; ++i) {
        int ql = qs[i][0];
        int qr = qs[i][1];
        int id = qs[i][2];
        int total = qr - ql + 1;
        while (l > ql) change(a[--l], +1);
        while (r < qr) change(a[++r], +1);
        while (l < ql) change(a[l++], -1);
        while (r > qr) change(a[r--], -1);
        
        vector<int> big;
        for (auto x : s) big.push_back(x);
        
        int start = 0;
        { // odd positions
            ll pref = 0, cnt = 0;
            for (int j = 1; j < K; ++j) {
                taken[j] = 0;
                if (cnt % 2 == 0) {
                    sol[id] += get(pref, j, total, (small[j] + 1) / 2);
                    taken[j] = (small[j] + 1) / 2;
                    pref += ((small[j] + 1) / 2) * 1LL * j;
                } else {
                    sol[id] += get(pref, j, total, small[j] / 2);
                    taken[j] = small[j] / 2;
                    pref += (small[j] / 2) * 1LL * j;
                }
                cnt += small[j];
            }
            for (int j = cnt % 2; j < big.size(); j += 2) {
                sol[id] += pref * (total - pref);
                pref += big[j];
            }
            start = cnt % 2;
        }
        { // even positions
            ll pref = 0, cnt = 0;
            for (int j = 1; j < K; ++j) { 
                sol[id] += get(pref + j, j, total, small[j] - taken[j]);
                cnt += small[j] - taken[j];
                pref += (small[j] - taken[j]) * 1LL * j;
            }
            for (int j = 1 - start; j < big.size(); j += 2) {
                pref += big[j];
                sol[id] += pref * (total - pref);
            }
        }
    }
    for (int i = 1; i <= q; ++i) cout << sol[i] << "\n";
}

• Subtask 1: 4 / 4
• Subtask 2: 10 / 10
• Subtask 3: 8 / 8
• Subtask 4: 16 / 16
• Subtask 5: 26 / 26
• Subtask 6: 36 / 36