data_structure/test/binary_trie.yuki2977.test.cpp

• View this file on GitHub
• Last update: 2024-12-07 22:13:21+09:00
• Problem: https://yukicoder.me/problems/no/2977

Depends on

• Binary trie (data_structure/binary_trie.hpp)

Code

#include "../binary_trie.hpp"
#define PROBLEM "https://yukicoder.me/problems/no/2977"

#include <iostream>
using namespace std;

int main() {
    cin.tie(nullptr), ios::sync_with_stdio(false);

    int N;
    long long K;
    cin >> N >> K;

    vector<int> A(N);
    for (auto &x : A) cin >> x;

    constexpr int D = 30;

    BinaryTrie<int> trie(D);
    for (int a : A) trie.insert(a);

    int lo = 0, hi = 1 << D; // [lo, hi)
    while (lo + 1 < hi) {
        const int mid = (lo + hi) / 2;

        long long cnt = 0;
        for (int a : A) cnt += trie.count_less_xor(a, mid);
        (cnt >= K * 2 + N ? hi : lo) = mid;
    }

    cout << lo << '\n';
}

#line 2 "data_structure/binary_trie.hpp"
#include <vector>

template <class Int, class Count = int> struct BinaryTrie {
    int maxD;
    std::vector<Count> deg, subtree_sum;
    std::vector<int> ch0, ch1, par;

    int _new_node(int id_par) {
        deg.emplace_back(Count());
        subtree_sum.emplace_back(Count());
        ch0.emplace_back(-1);
        ch1.emplace_back(-1);
        par.emplace_back(id_par);
        return (int)ch0.size() - 1;
    }

    BinaryTrie(int maxD = 0) : maxD(maxD) { _new_node(-1); }

    // Return index of x.
    // Create nodes to locate x if needed.
    int _get_create_index(Int x) {
        int now = 0;
        for (int d = maxD - 1; d >= 0; --d) {
            int nxt = ((x >> d) & 1) ? ch1[now] : ch0[now];
            if (nxt == -1) {
                nxt = _new_node(now);
                (((x >> d) & 1) ? ch1[now] : ch0[now]) = nxt;
            }
            now = nxt;
        }
        return now;
    }

    // Return node index of x.
    // Return -1 if x is not in trie.
    int _get_index(Int x) const {
        int now = 0;
        for (int d = maxD - 1; d >= 0; --d) {
            now = ((x >> d) & 1) ? ch1[now] : ch0[now];
            if (now == -1) return -1;
        }
        return now;
    }

    // insert x
    void insert(Int x, Count c = Count(1)) {
        int now = _get_create_index(x);
        deg[now] += c;
        while (now >= 0) subtree_sum[now] += c, now = par[now];
    }

    // delete x if exists
    void erase(Int x) {
        int now = _get_index(x);
        if (now >= 0 and deg[now] != 0) {
            Count r = deg[now];
            deg[now] = 0;
            while (now >= 0) subtree_sum[now] -= r, now = par[now];
        }
    }

    Count count(Int x) const {
        int now = _get_index(x);
        return now == -1 ? Count() : deg[now];
    }

    bool contains(Int x) const { return count(x) > Count(); }

    // min(y ^ x) for y in trie
    Int xor_min(Int x) const {
        Int ret = 0;
        int now = 0;
        if (!subtree_sum[now]) return -1;
        for (int d = maxD - 1; d >= 0; d--) {
            int y = ((x >> d) & 1) ? ch1[now] : ch0[now];
            if (y != -1 and subtree_sum[y]) {
                now = y;
            } else {
                ret += Int(1) << d, now = ch0[now] ^ ch1[now] ^ y;
            }
        }
        return ret;
    }

    // Count elements y such that x ^ y < thres
    Count count_less_xor(Int x, Int thres) const {
        Count ret = Count();
        int now = 0;

        for (int d = maxD - 1; d >= 0; d--) {
            if (now == -1) break;

            const bool bit_x = (x >> d) & 1;

            if ((thres >> d) & 1) {
                const int child = bit_x ? ch1[now] : ch0[now];
                if (child != -1) ret += subtree_sum[child];

                now = bit_x ? ch0[now] : ch1[now];
            } else {
                now = bit_x ? ch1[now] : ch0[now];
            }
        }

        return ret;
    }
};
#line 2 "data_structure/test/binary_trie.yuki2977.test.cpp"
#define PROBLEM "https://yukicoder.me/problems/no/2977"

#include <iostream>
using namespace std;

int main() {
    cin.tie(nullptr), ios::sync_with_stdio(false);

    int N;
    long long K;
    cin >> N >> K;

    vector<int> A(N);
    for (auto &x : A) cin >> x;

    constexpr int D = 30;

    BinaryTrie<int> trie(D);
    for (int a : A) trie.insert(a);

    int lo = 0, hi = 1 << D; // [lo, hi)
    while (lo + 1 < hi) {
        const int mid = (lo + hi) / 2;

        long long cnt = 0;
        for (int a : A) cnt += trie.count_less_xor(a, mid);
        (cnt >= K * 2 + N ? hi : lo) = mid;
    }

    cout << lo << '\n';
}

Back to top page