Binary trie
(data_structure/binary_trie.hpp)

• View this file on GitHub
• Last update: 2024-12-07 22:13:21+09:00
• Include: #include "data_structure/binary_trie.hpp"

非負整数の集合や多重集合に対する一部のクエリを効率的に行うためのデータ構造．

使用方法

using Key = int;
using Count = int;
const int D = 30;  // Key の桁数

BinaryTrie<Key, Count> trie(D);

for (int a : A) trie.insert(a);

Key t;
Count n = trie.count_less_xor(a, t);  // a ^ x < t を満たす x （x は現在存在する値）を数える

Key v = bt.xor_min(t);  // t ^ x （x は現在存在する値）の最小値を求める

問題例

• Library Checker: Set Xor-Min
• No.2977 Kth Xor Pair - yukicoder

Verified with

• data_structure/test/binary_trie.test.cpp
• data_structure/test/binary_trie.yuki2977.test.cpp

Code

#pragma once
#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/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;
    }
};

Back to top page