Submission #798123

#	Time	Username	Problem	Language	Result	Execution time	Memory
798123		finn__	Diversity (CEOI21_diversity)	C++17	100 / 100	6451 ms	6532 KiB

diversity

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

template <size_t N>
struct OrSegmentTree
{
    bitset<2 * N> t;

    bool operator[](size_t i) { return t[N + i]; }
    void reset() { t.reset(); }

    void set(size_t i, bool x)
    {
        t[i += N] = x;
        while (i >>= 1)
            t[i] = t[2 * i] | t[2 * i + 1];
    }

    int64_t previous_nonzero(size_t i)
    {
        i += N;
        do
        {
            if ((i & 1) && t[i - 1])
            {
                --i;
                while (i < N)
                    i = 2 * i + t[2 * i + 1];
                return i - N;
            }
        } while ((i >>= 1) > 1);
        return -1;
    }

    int64_t next_nonzero(size_t i)
    {
        i += N;
        do
        {
            if (!(i & 1) && t[i + 1])
            {
                ++i;
                while (i < N)
                    i = 2 * i + !t[2 * i];
                return i - N;
            }
        } while ((i >>= 1) > 1);
        return N;
    }
};

constexpr size_t N = 300000, K = 2400;

uint32_t a[N], freq[N], freq_of_freq[N];
uint64_t ans[N];
tuple<size_t, size_t, size_t> queries[N];
OrSegmentTree<N> tree;

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

inline uint64_t range_gauss(uint64_t a, uint64_t b)
{
    return gauss(b) - (a ? gauss(a - 1) : 0);
}

inline uint64_t sum_of_squares(uint64_t n) { return (n * (n + 1) * ((n << 1) + 1)) / 6; }

inline uint64_t range_cost(uint64_t l0, uint64_t y, uint64_t z, uint64_t n)
{
    return -range_gauss(l0, l0 + z * y - 1) - z * l0 * l0 - 2 * l0 * y * gauss(z - 1) -
           y * y * sum_of_squares(z - 1) + z * n * (l0 + y) + y * n * gauss(z - 1);
}

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

    size_t n, q;
    cin >> n >> q;

    for (size_t i = 0; i < n; ++i)
        cin >> a[i], --a[i];
    for (size_t i = 0; i < q; ++i)
        cin >> get<0>(queries[i]) >> get<1>(queries[i]), --get<0>(queries[i]), --get<1>(queries[i]),
            get<2>(queries[i]) = i;
    sort(queries, queries + q, [](auto const &a, auto const &b)
         { return get<0>(a) / K == get<0>(b) / K ? get<1>(a) > get<1>(b) : get<0>(a) < get<0>(b); });

    size_t i = 0, j = 0;
    freq[a[0]] = 1;
    freq_of_freq[1] = 1;
    tree.set(freq[a[0]], 1);

    for (size_t k = 0; k < q; ++k)
    {
        auto const [u, v, qi] = queries[k];
        while (j < v)
        {
            ++j;
            freq[a[j]]++;
            freq_of_freq[freq[a[j]] - 1]--;
            freq_of_freq[freq[a[j]]]++;
            if (!freq_of_freq[freq[a[j]] - 1])
                tree.set(freq[a[j]] - 1, 0);
            if (freq_of_freq[freq[a[j]]] == 1)
                tree.set(freq[a[j]], 1);
        }
        while (i > u)
        {
            --i;
            freq[a[i]]++;
            freq_of_freq[freq[a[i]] - 1]--;
            freq_of_freq[freq[a[i]]]++;
            if (!freq_of_freq[freq[a[i]] - 1])
                tree.set(freq[a[i]] - 1, 0);
            if (freq_of_freq[freq[a[i]]] == 1)
                tree.set(freq[a[i]], 1);
        }
        while (j > v)
        {
            freq_of_freq[freq[a[j]]]--;
            freq_of_freq[freq[a[j]] - 1]++;
            if (!freq_of_freq[freq[a[j]]])
                tree.set(freq[a[j]], 0);
            if (freq_of_freq[freq[a[j]] - 1] == 1)
                tree.set(freq[a[j]] - 1, 1);
            freq[a[j]]--;
            --j;
        }
        while (i < u)
        {
            freq_of_freq[freq[a[i]]]--;
            freq_of_freq[freq[a[i]] - 1]++;
            if (!freq_of_freq[freq[a[i]]])
                tree.set(freq[a[i]], 0);
            if (freq_of_freq[freq[a[i]] - 1] == 1)
                tree.set(freq[a[i]] - 1, 1);
            freq[a[i]]--;
            ++i;
        }

        vector<pair<uint64_t, uint64_t>> y;
        size_t i = tree.next_nonzero(0);
        while (i < N)
        {
            y.emplace_back(i, freq_of_freq[i]);
            i = tree.next_nonzero(i);
        }

        uint64_t lsum = 0, rsum = 0, cost = 0;
        for (auto const &[freq, num] : y)
            if (freq)
            {
                size_t num_left, num_right;
                if (lsum < rsum)
                    num_right = num / 2, num_left = num - num_right;
                else
                    num_left = num / 2, num_right = num - num_left;
                cost += range_cost(lsum, freq, num_left, v - u + 1);
                cost += range_cost((v - u + 1) - rsum - num_right * freq, freq, num_right, v - u + 1);
                lsum += num_left * freq;
                rsum += num_right * freq;
            }
        ans[qi] = cost;
    }

    for (size_t i = 0; i < q; ++i)
        cout << ans[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