Splay 板子

__ryp__ · 2024-02-16 10:40:03 · 个人记录

挺简洁的，效率也不错。去掉类模板 78 行。

template<typename T, size_t N> class SplayTree
{
private:
    T val[N];
    int son[N][2], fa[N], siz[N], cntq[N], cnt = 1, rt = 0;

    void maintain (int x)
    { siz[x] = siz[son[x][0]] + siz[son[x][1]] + cntq[x]; }
    int get (int x) const { return son[fa[x]][1] == x; }
    int compare (int f, T k) const { return k > val[f]; }
    void rotate (int);
    int splay (int f, int t = 0);
    int search (T k);
    int insert_at (int f, T k) {
        cntq[cnt] = siz[cnt] = 1, val[cnt] = k, fa[cnt] = f;
        if (f) son[f][compare (f, k)] = cnt;
        return cnt++;
    }
    int merge (int x, int y);
public: 
    int insert (T k);
    int remove (T k);
    int qrank (T k) { return siz[son[search (k)][0]] + cntq[rt] * (val[rt] < k); }
    T kth (int k);
    T qpre (T k);
    T qnext (T k);
};

template<typename T, size_t N>
void SplayTree<T, N>::rotate (int x)
{
    int y = fa[x], z = fa[y], p = get (x);
    son[y][p] = son[x][1 ^ p];
    if (son[x][1 ^ p]) fa[son[x][p ^ 1]] = y;
    son[x][p ^ 1] = y, fa[y] =  x, fa[x] = z;
    if (z) son[z][y == son[z][1]] = x;
    maintain (y), maintain (x);
}

template<typename T, size_t N>
int SplayTree<T, N>::splay (int x, int t)
{
    for (int f = fa[x]; (f = fa[x]) != t; rotate (x))
        if (fa[f] != t) rotate (get (x) == get (f) ? f : x);
    return t ? x : rt = x;
}

template<typename T, size_t N>
int SplayTree<T, N>::insert (T k)
{
    int f = rt, last = 0;
    if (!f) return rt = insert_at (0, k);
    while (f && val[f] != k) last = f, f = son[f][compare (f, k)];
    if (f && val[f] == k) return cntq[splay (f)]++;
    else return f = insert_at (last, k), maintain (last), splay (f);
}

template<typename T, size_t N>
int SplayTree<T, N>::merge (int x, int y)
{
    int f = x;
    if (!x || !y) return x + y;
    while (son[f][1]) f = son[f][1];
    splay (f, fa[x]);
    if (y) fa[y] = f;
    son[f][1] = y;
    return f;
}

template<typename T, size_t N>
int SplayTree<T, N>::remove (T k)
{
    int f = search (k), x;
    if (cntq[f] > 1) return cntq[f]--, f;
    if ((x = merge (son[f][0], son[f][1]))) fa[x] = 0;
    return rt = x;
}

template<typename T, size_t N>
T SplayTree<T, N>::kth (int k)
{
    int f = rt;
    while (f)
        if (k <= siz[son[f][0]]) f = son[f][0];
        else if (k <= siz[son[f][0]] + cntq[f]) return val[splay (f)];
        else k -= siz[son[f][0]] + cntq[f], f = son[f][1];
    return 0;
}

template<typename T, size_t N>
T SplayTree<T, N>::qpre (T k)
{
    int f = son[search (k)][0];
    if (val[rt] < k) return val[rt];
    while (son[f][1]) f = son[f][1];
    return val[splay (f)];
}

template<typename T, size_t N>
T SplayTree<T, N>::qnext (T k)
{
    int f = son[search (k)][1];
    if (val[rt] > k) return val[rt];
    while (son[f][0]) f = son[f][0];
    return val[splay (f)];
}

template<typename T, size_t N>
int SplayTree<T, N>::search (T k)
{
    int f = rt;
    if (!f) return 0;
    while (son[f][compare (f, k)] && val[f] != k)
        f = son[f][compare (f, k)];
    return splay (f);
}