Submission #520490

#	Time	Username	Problem	Language	Result	Execution time	Memory
520490		KoD	Diversity (CEOI21_diversity)	C++17	100 / 100	4338 ms	8264 KiB

diversity

#include <bits/stdc++.h>

using std::vector;
using std::array;
using std::pair;
using std::tuple;

using i64 = std::int64_t;

constexpr int B = 700;
constexpr int MAX = 300000;

constexpr i64 sum_k(const int n) {
    return (i64)n * (n + 1) / 2;
}
 
constexpr i64 sum_k2(const int n) {
    return (i64)n * (n + 1) * (2 * n + 1) / 6;
}

int main() {
    std::ios_base::sync_with_stdio(false);
    std::cin.tie(nullptr);
    int N, Q;
    std::cin >> N >> Q;
    vector<int> A(N);
    for (auto& x : A) {
        std::cin >> x;
        x -= 1;
    }
    vector<vector<int>> query((N + B - 1) / B);
    vector<int> L(Q), R(Q);
    for (int i = 0; i < Q; ++i) {
        std::cin >> L[i] >> R[i];
        L[i] -= 1;
        query[L[i] / B].push_back(i);
    }
    vector<i64> ans(Q);
    vector<int> count(MAX), block(N + 1), nonzero;
    const auto change = [&](const int x, const int c) {
        block[count[x]] -= 1;
        count[x] += c;
        if (count[x] != 0 and block[count[x]] == 0) {
            nonzero.push_back(count[x]);
        }
        block[count[x]] += 1;
    };
    for (auto& qidx : query) {
        if (qidx.empty()) {
            continue;
        }
        std::sort(qidx.begin(), qidx.end(), [&](const int i, const int j) {
            return R[i] < R[j]; 
        });
        std::fill(count.begin(), count.end(), 0);
        std::fill(block.begin(), block.end(), 0);
        nonzero.clear();
        int l = L[qidx[0]], r = L[qidx[0]];
        for (const int q : qidx) {
            while (r < R[q]) {
                change(A[r++], 1);
            }
            while (l > L[q]) {
                change(A[--l], 1);
            }
            while (l < L[q]) {
                change(A[l++], -1);
            }
            std::sort(nonzero.begin(), nonzero.end());
            nonzero.erase(std::unique(nonzero.begin(), nonzero.end()), nonzero.end());
            const int len = R[q] - L[q];
            ans[q] = sum_k(len);
            const auto add_to = [&](const i64 s, const int a, const int n) {
                ans[q] += (n * s + sum_k(n) * a) * len;
                ans[q] -= n * s * s + 2 * s * sum_k(n) * a + sum_k2(n) * a * a;
            };
            int left = 0, right = 0;
            vector<int> next;
            for (const int a : nonzero) {
                const int b = block[a];
                if (b == 0) {
                    continue;
                }
                next.push_back(a);
                if (left > right) {
                    std::swap(left, right);
                }
                const int need = std::min(b, (right - left + a - 1) / a);
                const int to_left = need + (b - need) / 2;
                const int to_right = b - to_left;
                add_to(left, a, to_left);
                add_to(right, a, to_right);
                left += a * to_left;
                right += a * to_right;
            }
            ans[q] -= (i64)left * (len - left);
            nonzero = std::move(next);
        }
    }
    for (const i64 x : ans) {
        std::cout << x << '\n';
    }
    return 0;
}

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