【个人笔记】BST二叉搜索树模板

【个人笔记】BST（Binary Search Tree）二叉搜索树代码模板

#include <bits/stdc++.h>
#define endl '\n'

using namespace std;

const int INF = 0x7fffffff;
const double EPS = 1e-8;
const int N = 1e4 + 5;

struct node {
    int left;
    int right;
    int value;
    int cnt;
    int siz;
} tree[N];

int q, ncnt;

void add(int x, int root) {
    tree[root].siz++;
    if (tree[root].value == x) {
        tree[root].cnt++;
        return;
    }
    if (x < tree[root].value) {
        if (tree[root].left) {
            add(x, tree[root].left);
        } else {
            tree[++ncnt].value = x;
            tree[ncnt].siz = tree[ncnt].cnt = 1;
            tree[root].left = ncnt;
        }
    } else {
        if (tree[root].right) {
            add(x, tree[root].right);
        } else {
            tree[++ncnt].value = x;
            tree[ncnt].siz = tree[ncnt].cnt = 1;
            tree[root].right = ncnt;
        }
    }
}

int pre(int x, int root, int ans) {
    if (tree[root].value >= x) {
        if (tree[root].left) {
            return pre(x, tree[root].left, ans);
        } else {
            return ans;
        }
    } else {
        if (!tree[root].right) {
            return (tree[root].value < x) ? tree[root].value : ans;
        }
        if (tree[root].cnt != 0) {
            return pre(x, tree[root].right, tree[root].value);
        } else {
            return pre(x, tree[root].right, ans);
        }
    }
}

int post(int x, int root, int ans) {
    if (tree[root].value <= x) {
        if (tree[root].right)
            return post(x, tree[root].right, ans);
        else
            return ans;
    } else {
        if (!tree[root].left) {
            return (tree[root].value > x) ? tree[root].value : ans;
        }
        if (tree[root].cnt != 0) {
            return post(x, tree[root].left, tree[root].value);
        } else {
            return post(x, tree[root].left, ans);
        }
    }
}

int rk(int x, int root) {
    if (root == 0)
        return 0;
    if (x == tree[root].value)
        return tree[tree[root].left].siz;
    if (x < tree[root].value)
        return rk(x, tree[root].left);
    else
        return rk(x, tree[root].right) + tree[tree[root].left].siz +
               tree[root].cnt;
}

int kth(int x, int root) {
    if (root == 0)
        return INF;
    if (tree[tree[root].left].siz >= x) {
        return kth(x, tree[root].left);
    }
    if (tree[tree[root].left].siz + tree[root].cnt >= x) {
        return tree[root].value;
    }
    return kth(x - tree[tree[root].left].siz - tree[root].cnt,
               tree[root].right);
}

int main() {
    cin >> q;
    while (q--) {
        int op, x;
        cin >> op >> x;
        switch (op) {
            case 1:
                cout << rk(x, 1) + 1 << endl;
                break;
            case 2:
                cout << kth(x, 1) << endl;
                break;
            case 3:
                cout << pre(x, 1, -INF) << endl;
                break;
            case 4:
                cout << post(x, 1, INF) << endl;
                break;
            case 5:
                if (ncnt == 0) {
                    ncnt++;
                    tree[ncnt].value = x;
                    tree[ncnt].cnt = tree[ncnt].siz = 1;
                } else
                    add(x, 1);
                break;
            default:
                break;
        }
    }
    return 0;
}