Persistent Randomized Binary Search Tree (set)
(DataStructure/persistentrbstset.hpp)

• View this file on GitHub
• Last update: 2024-04-26 03:18:17+09:00
• Include: #include "DataStructure/persistentrbstset.hpp"

Depends on

• Random (Utility/random.hpp)

Verified with

• Verify/LC_range_kth_smallest-2.test.cpp

Code

#pragma once
#include "Utility/random.hpp"


template <typename T, int LIM = 10101010> struct PRBSTset {
    struct Node {
        Node *lp = nullptr, *rp = nullptr;
        int sz = 1;
        T val;
        Node() {}
        void apply() {
            sz = 1;
            if (lp)
                sz += lp->sz;
            if (rp)
                sz += rp->sz;
        }
    };
    using np = Node *;
    Node buf[LIM];
    int pos = 0;
    int sz(np root) { return root ? root->sz : 0; }
    np merge(np L, np R) {
        if (!L)
            return R;
        if (!R)
            return L;
        if (Random::uniform() * (sz(L) + sz(R)) < sz(L)) {
            auto rb = merge(L->rp, R);
            np ret = make(L->val, L->lp, rb);
            return ret;
        } else {
            auto lb = merge(L, R->lp);
            np ret = make(R->val, lb, R->rp);
            return ret;
        }
    }
    array<np, 2> split(np root, int k) {
        if (k <= 0)
            return {nullptr, root};
        if (k >= sz(root))
            return {root, nullptr};
        if (k <= sz(root->lp)) {
            auto [L, lb] = split(root->lp, k);
            np R = make(root->val, lb, root->rp);
            return {L, R};
        } else {
            auto [rb, R] = split(root->rp, k - 1 - sz(root->lp));
            np L = make(root->val, root->lp, rb);
            return {L, R};
        }
    }
    bool find(np root, T v) {
        if (!root)
            return false;
        if (root->val == v)
            return true;
        else if (root->val > v)
            return find(root->lp, v);
        else
            return find(root->rp, v);
    }
    int lower_bound(np root, T v) {
        if (!root)
            return 0;
        if (root->val > v)
            return lower_bound(root->lp, v);
        else
            return sz(root->lp) + 1 + lower_bound(root->rp, v);
    }
    int upper_bound(np root, T v) {
        if (!root)
            return 0;
        if (root->val >= v)
            return upper_bound(root->lp, v);
        else
            return sz(root->lp) + 1 + upper_bound(root->rp, v);
    }
    np make(T v, np L = nullptr, np R = nullptr) {
        np ret = &buf[pos++];
        ret->val = v;
        ret->lp = L;
        ret->rp = R;
        ret->apply();
        return ret;
    }
    void dfs(np root, vector<T> &a) {
        if (!root)
            return;
        dfs(root->lp, a);
        a.push_back(root->val);
        dfs(root->rp, a);
    }
    np rebuild(np root) {
        if (pos < LIM * .95)
            return root;
        vector<T> tmp;
        dfs(root, tmp);
        return build(tmp);
    }
    np insert(np root, T v) {
        int k = lower_bound(root, v);
        auto [L, R] = split(root, k);
        return merge(merge(L, make(v)), R);
    }
    np erase(np root, T v) {
        int k = lower_bound(root, v);
        auto [L, rb] = split(root, k);
        auto [tmp, R] = split(rb, 1);
        return merge(L, R);
    }
    T kth_elem(np root, int k) {
        assert(k <= 0 and k < sz(root));
        if (sz(root->lp) == k)
            return root->val;
        else if (sz(root->lp) > k)
            return kth_elem(root->lp, k);
        else
            return kth_elem(root->rp, k - 1 - sz(root->lp));
    }
    np build(vector<T> &a) {
        np root = nullptr;
        for (auto &x : a)
            root = merge(root, make(x));
        return root;
    }
    void dump(np root) {
        if (!root)
            return;
        dump(root->lp);
        cerr << root->val << '\n';
        dump(root->rp);
    }
};

/**
 * @brief Persistent Randomized Binary Search Tree (set)
 */

#line 2 "Utility/random.hpp"

namespace Random {
mt19937_64 randgen(chrono::steady_clock::now().time_since_epoch().count());
using u64 = unsigned long long;
u64 get() {
    return randgen();
}
template <typename T> T get(T L) { // [0,L]

    return get() % (L + 1);
}
template <typename T> T get(T L, T R) { // [L,R]

    return get(R - L) + L;
}
double uniform() {
    return double(get(1000000000)) / 1000000000;
}
string str(int n) {
    string ret;
    rep(i, 0, n) ret += get('a', 'z');
    return ret;
}
template <typename Iter> void shuffle(Iter first, Iter last) {
    if (first == last)
        return;
    int len = 1;
    for (auto it = first + 1; it != last; it++) {
        len++;
        int j = get(0, len - 1);
        if (j != len - 1)
            iter_swap(it, first + j);
    }
}
template <typename T> vector<T> select(int n, T L, T R) { // [L,R]

    if (n * 2 >= R - L + 1) {
        vector<T> ret(R - L + 1);
        iota(ALL(ret), L);
        shuffle(ALL(ret));
        ret.resize(n);
        return ret;
    } else {
        unordered_set<T> used;
        vector<T> ret;
        while (SZ(used) < n) {
            T x = get(L, R);
            if (!used.count(x)) {
                used.insert(x);
                ret.push_back(x);
            }
        }
        return ret;
    }
}

void relabel(int n, vector<pair<int, int>> &es) {
    shuffle(ALL(es));
    vector<int> ord(n);
    iota(ALL(ord), 0);
    shuffle(ALL(ord));
    for (auto &[u, v] : es)
        u = ord[u], v = ord[v];
}
template <bool directed, bool simple> vector<pair<int, int>> genGraph(int n) {
    vector<pair<int, int>> cand, es;
    rep(u, 0, n) rep(v, 0, n) {
        if (simple and u == v)
            continue;
        if (!directed and u > v)
            continue;
        cand.push_back({u, v});
    }
    int m = get(SZ(cand));
    vector<int> ord;
    if (simple)
        ord = select(m, 0, SZ(cand) - 1);
    else {
        rep(_, 0, m) ord.push_back(get(SZ(cand) - 1));
    }
    for (auto &i : ord)
        es.push_back(cand[i]);
    relabel(n, es);
    return es;
}
vector<pair<int, int>> genTree(int n) {
    vector<pair<int, int>> es;
    rep(i, 1, n) es.push_back({get(i - 1), i});
    relabel(n, es);
    return es;
}
}; // namespace Random


/**
 * @brief Random
 */
#line 3 "DataStructure/persistentrbstset.hpp"

template <typename T, int LIM = 10101010> struct PRBSTset {
    struct Node {
        Node *lp = nullptr, *rp = nullptr;
        int sz = 1;
        T val;
        Node() {}
        void apply() {
            sz = 1;
            if (lp)
                sz += lp->sz;
            if (rp)
                sz += rp->sz;
        }
    };
    using np = Node *;
    Node buf[LIM];
    int pos = 0;
    int sz(np root) { return root ? root->sz : 0; }
    np merge(np L, np R) {
        if (!L)
            return R;
        if (!R)
            return L;
        if (Random::uniform() * (sz(L) + sz(R)) < sz(L)) {
            auto rb = merge(L->rp, R);
            np ret = make(L->val, L->lp, rb);
            return ret;
        } else {
            auto lb = merge(L, R->lp);
            np ret = make(R->val, lb, R->rp);
            return ret;
        }
    }
    array<np, 2> split(np root, int k) {
        if (k <= 0)
            return {nullptr, root};
        if (k >= sz(root))
            return {root, nullptr};
        if (k <= sz(root->lp)) {
            auto [L, lb] = split(root->lp, k);
            np R = make(root->val, lb, root->rp);
            return {L, R};
        } else {
            auto [rb, R] = split(root->rp, k - 1 - sz(root->lp));
            np L = make(root->val, root->lp, rb);
            return {L, R};
        }
    }
    bool find(np root, T v) {
        if (!root)
            return false;
        if (root->val == v)
            return true;
        else if (root->val > v)
            return find(root->lp, v);
        else
            return find(root->rp, v);
    }
    int lower_bound(np root, T v) {
        if (!root)
            return 0;
        if (root->val > v)
            return lower_bound(root->lp, v);
        else
            return sz(root->lp) + 1 + lower_bound(root->rp, v);
    }
    int upper_bound(np root, T v) {
        if (!root)
            return 0;
        if (root->val >= v)
            return upper_bound(root->lp, v);
        else
            return sz(root->lp) + 1 + upper_bound(root->rp, v);
    }
    np make(T v, np L = nullptr, np R = nullptr) {
        np ret = &buf[pos++];
        ret->val = v;
        ret->lp = L;
        ret->rp = R;
        ret->apply();
        return ret;
    }
    void dfs(np root, vector<T> &a) {
        if (!root)
            return;
        dfs(root->lp, a);
        a.push_back(root->val);
        dfs(root->rp, a);
    }
    np rebuild(np root) {
        if (pos < LIM * .95)
            return root;
        vector<T> tmp;
        dfs(root, tmp);
        return build(tmp);
    }
    np insert(np root, T v) {
        int k = lower_bound(root, v);
        auto [L, R] = split(root, k);
        return merge(merge(L, make(v)), R);
    }
    np erase(np root, T v) {
        int k = lower_bound(root, v);
        auto [L, rb] = split(root, k);
        auto [tmp, R] = split(rb, 1);
        return merge(L, R);
    }
    T kth_elem(np root, int k) {
        assert(k <= 0 and k < sz(root));
        if (sz(root->lp) == k)
            return root->val;
        else if (sz(root->lp) > k)
            return kth_elem(root->lp, k);
        else
            return kth_elem(root->rp, k - 1 - sz(root->lp));
    }
    np build(vector<T> &a) {
        np root = nullptr;
        for (auto &x : a)
            root = merge(root, make(x));
        return root;
    }
    void dump(np root) {
        if (!root)
            return;
        dump(root->lp);
        cerr << root->val << '\n';
        dump(root->rp);
    }
};

/**
 * @brief Persistent Randomized Binary Search Tree (set)
 */

Back to top page