CP Notebook

← all categories

Splay Implicit 53ebabf0

Implicit (positional) splay tree: array-as-balanced-BST via splaying. Supports insert-at-position, erase-at-position, range reverse, and an "op" block (aggregate + lazy tag) you swap for the problem at hand — default here is an affine range update (x -> mul*x + add) + range-sum, a richer example than treap-implicit.h's additive one; same op-block pattern (see treap-implicit.h). T needs operator+, operator*, and default-constructed 0/1. Uses two internal sentinel positions. Multi-page exception to the line budget.

Time: O(\log n) amortized per operation. tested

content/data-structures/splay-implicit.h

template <typename T>
struct SplayImplicit {
    struct Node {
        // ---- op: edit for a different aggregate/lazy tag (see treap-implicit.h) ----
        T val, agg;
        T mul = T(1), add = T();  // affine lazy: x -> mul*x + add
        // ---- end op ----
        int p = -1, ch[2] = {-1, -1}, sz = 1, rev = 0;
    };
    vector<Node> pool;
    int root = -1, n = 0;  // n = logical size, excludes the 2 sentinels

    void apply(int u, T mul, T add) {  // op: push an affine tag onto u
        pool[u].val = mul * pool[u].val + add;
        pool[u].agg = mul * pool[u].agg + add * pool[u].sz;
        pool[u].mul = mul * pool[u].mul;
        pool[u].add = mul * pool[u].add + add;
    }

    int getSz(int u) { return u == -1 ? 0 : pool[u].sz; }
    T getAgg(int u) { return u == -1 ? T() : pool[u].agg; }
    int make(T x) { pool.push_back({x, x}); return (int)pool.size() - 1; }

    void pull(int u) {
        pool[u].sz = 1 + getSz(pool[u].ch[0]) + getSz(pool[u].ch[1]);
        pool[u].agg = getAgg(pool[u].ch[0]) + pool[u].val + getAgg(pool[u].ch[1]);
    }

    void push(int u) {
        if (pool[u].rev) {
            swap(pool[u].ch[0], pool[u].ch[1]);
            if (pool[u].ch[0] != -1) pool[pool[u].ch[0]].rev ^= 1;
            if (pool[u].ch[1] != -1) pool[pool[u].ch[1]].rev ^= 1;
            pool[u].rev = 0;
        }
        if (pool[u].mul != T(1) || pool[u].add != T()) {
            if (pool[u].ch[0] != -1) apply(pool[u].ch[0], pool[u].mul, pool[u].add);
            if (pool[u].ch[1] != -1) apply(pool[u].ch[1], pool[u].mul, pool[u].add);
            pool[u].mul = T(1); pool[u].add = T();
        }
    }

    int dir(int u) { return pool[pool[u].p].ch[1] == u; }

    void rotate(int u) {
        int f = pool[u].p, g = pool[f].p, d = dir(u), w = pool[u].ch[d ^ 1];
        pool[f].ch[d] = w;
        if (w != -1) pool[w].p = f;
        pool[u].ch[d ^ 1] = f;
        pool[f].p = u;
        pool[u].p = g;
        if (g != -1) {
            if (pool[g].ch[0] == f) pool[g].ch[0] = u;
            else if (pool[g].ch[1] == f) pool[g].ch[1] = u;
        }
        pull(f); pull(u);
    }

    void pushPath(int u) {
        if (pool[u].p != -1) pushPath(pool[u].p);
        push(u);
    }

    void splay(int u, int g = -1) {
        pushPath(u);
        while (pool[u].p != g) {
            int f = pool[u].p, gg = pool[f].p;
            if (gg != g) {
                if (dir(u) == dir(f)) rotate(f);
                else rotate(u);
            }
            rotate(u);
        }
        if (g == -1) root = u;
    }

    int findKth(int k, int g = -1) {  // 0-indexed over the internal (sentinel-padded) array
        int u = root;
        while (u != -1) {
            push(u);
            int szL = getSz(pool[u].ch[0]);
            if (k < szL) u = pool[u].ch[0];
            else if (k == szL) { splay(u, g); return u; }
            else { k -= szL + 1; u = pool[u].ch[1]; }
        }
        return -1;
    }

    int buildRec(vector<T> &a, int l, int r) {
        if (l > r) return -1;
        int m = (l + r) / 2;
        int u = make(a[m]);
        pool[u].ch[0] = buildRec(a, l, m - 1);
        pool[u].ch[1] = buildRec(a, m + 1, r);
        if (pool[u].ch[0] != -1) pool[pool[u].ch[0]].p = u;
        if (pool[u].ch[1] != -1) pool[pool[u].ch[1]].p = u;
        pull(u);
        return u;
    }

    void build(vector<T> &a) {  // adds 2 internal sentinels around a
        n = (int)a.size();
        vector<T> b(n + 2);
        for (int i = 0; i < n; i++) b[i + 1] = a[i];
        root = buildRec(b, 0, n + 1);
    }

    int isolate(int lo, int hi) {  // 0-indexed inclusive; returns subtree root for [lo, hi]
        int L = findKth(lo);
        int R = findKth(hi + 2, L);
        return pool[R].ch[0];
    }

    void insert(int pos, T x) {  // 0-indexed, shifts [pos, n) right
        int L = findKth(pos);
        int R = findKth(pos + 1, L);
        int m = make(x);
        pool[m].p = R;
        pool[R].ch[0] = m;
        pull(R); pull(L);
        n++;
    }

    void erase(int pos) {  // 0-indexed
        int L = findKth(pos);
        int R = findKth(pos + 2, L);
        pool[R].ch[0] = -1;
        pull(R); pull(L);
        n--;
    }

    void reverse(int lo, int hi) { pool[isolate(lo, hi)].rev ^= 1; }
    void update(int lo, int hi, T mul, T add) { apply(isolate(lo, hi), mul, add); }
    T query(int lo, int hi) { return getAgg(isolate(lo, hi)); }
};