Submission #789773

#	Time	Username	Problem	Language	Result	Execution time	Memory
789773		math_rabbit_1028	Diversity (CEOI21_diversity)	C++14	100 / 100	6521 ms	3840 KiB

diversity

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
int n, q, b;
int arr[303030];
ll cnt[303030], cnt2[303030], sum[303030], xsum[303030], ans[50505];
pair<ll, ll> reord[303030];
unordered_set<int> S;

const int m = 300000;

struct query{
    int idx, lt, rt;
} Q[50505];

bool compare(query p, query q) {
    if (p.lt/b == q.lt/b) return p.rt < q.rt;
    return (p.lt/b) < (q.lt/b);
}

void add(int i, int diff) {
    if (cnt2[cnt[i]] == 1) S.erase(cnt[i]);
    cnt2[cnt[i]]--;
    cnt[i] += diff;
    if (cnt2[cnt[i]] == 0) S.insert(cnt[i]);
    cnt2[cnt[i]]++;
}

ll SumOfSquare(ll x, ll y, ll a, ll k) { // \sum_{i=0}^{k-1} {(x+ia)^2 + (y+ia)^2}
    return k * x * x + a * x * k * (k - 1) + a * a * k * (k - 1) * (2 * k - 1) / 6
        + k * y * y + a * y * k * (k - 1) + a * a * k * (k - 1) * (2 * k - 1) / 6;
}

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

    cin >> n >> q; b = (int)sqrt(n);
    for (int i = 1; i <= n; i++) cin >> arr[i];

    for (int i = 1; i <= q; i++) cin >> Q[i].lt >> Q[i].rt;
    for (int i = 1; i <= q; i++) Q[i].idx = i;

    sort(Q + 1, Q + q + 1, compare);

    Q[0].lt = 1; Q[0].rt = 0;
    for (int i = 1; i <= q; i++) {
        if (Q[i - 1].rt < Q[i].rt) for (int j = Q[i - 1].rt + 1; j <= Q[i].rt; j++) add(arr[j], 1);
        else for (int j = Q[i].rt + 1; j <= Q[i - 1].rt; j++) add(arr[j], -1);

        if (Q[i - 1].lt > Q[i].lt) for (int j = Q[i].lt; j < Q[i - 1].lt; j++) add(arr[j], 1);
        else for (int j = Q[i - 1].lt; j < Q[i].lt; j++) add(arr[j], -1);

        ll len = Q[i].rt - Q[i].lt + 1;
        ll full_div = 0;
        vector< pair<int, ll> > v;
        for (auto iter = S.begin(); iter != S.end(); iter++) {
            v.emplace_back(*iter, cnt2[*iter]);
            full_div += cnt2[*iter];
        }

        sort(v.begin(), v.end());

        ll res = full_div * len * (len + 1) - len * (full_div - 1);
        ll l = 0, r = 0;
        for (int i = 0; i < v.size(); i++) {
            ll a = v[i].first, k = v[i].second;
            ll sm = k/2, lr = k - k/2;
            res -= SumOfSquare(l, len - l - a * sm, a, sm);
            res -= SumOfSquare(r, len - r - a * lr, a, lr);
            l = l + a * sm;
            r = r + a * lr;
            if (k & 1) swap(l, r);
        }
        ans[Q[i].idx] = res/2;
    }

    for (int i = 1; i <= q; i++) cout << ans[i] << "\n";

    return 0;
}

Compilation message (stderr)

diversity.cpp: In function 'int main()':
diversity.cpp:66:27: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   66 |         for (int i = 0; i < v.size(); i++) {
      |                         ~~^~~~~~~~~~

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