Untitled
unknown
c_cpp
a month ago
3.8 kB
0
Indexable
Never
#include <iostream> #include <algorithm> using namespace std; struct node { int key, bf, h, sz, lsz; node *l, *r, *pa; void update(); node() {} node(int _key) : key(_key), bf(0), h(0), sz(1), lsz(1), pa(NULL), l(NULL), r(NULL) {} } *root; void node::update() { int lh = (l ? l->h : -1); int rh = (r ? r->h : -1); int tlsz = (l ? l->sz : 0); int rsz = (r ? r->sz : 0); sz = tlsz + rsz + 1; lsz = tlsz + 1; h = max(lh, rh) + 1; bf = lh - rh; } node *find_min(node *x) { while (x->l) x = x->l; return x; } void find_kth(node *x, int k) { if (x->sz < k) { cout << "-1\n"; return; } if (x->lsz == k) { cout << x->key << "\n"; return; } else if (x->lsz > k) { x = x->l; find_kth(x, k); } else if (x->lsz < k) { k -= x->lsz; x = x->r; find_kth(x, k); } } void trans(node *x, node *y) { if (!x->pa) root = y; else if (x->pa->l == x) x->pa->l = y; else x->pa->r = y; if (y) y->pa = x->pa; } void rrot(node *x) { node *left_child = x->l; trans(x, left_child); x->l = left_child->r; if (left_child->r) left_child->r->pa = x; left_child->r = x; x->pa = left_child; x->update(); left_child->update(); } void lrot(node *x) { node *right_child = x->r; trans(x, right_child); x->r = right_child->l; if (right_child->l) right_child->l->pa = x; right_child->l = x; x->pa = right_child; x->update(); right_child->update(); } void insert(node *&x, node *z) { if (x == NULL) { x = z; return; } if (z->key == x->key) return; if (x->key < z->key) { if (x->r) insert(x->r, z); else x->r = z, z->pa = x; } else { if (x->l) insert(x->l, z); else x->l = z, z->pa = x; } x->update(); if (x->bf == 2) { if (x->l->bf == 1) rrot(x); else if (x->l->bf == -1) lrot(x->l), rrot(x); } else if (x->bf == -2) { if (x->r->bf == -1) lrot(x); else if (x->r->bf == 1) rrot(x->r), lrot(x); } } void del(node *x, int key) { if (x == NULL) return; if (key < x->key) del(x->l, key); else if (key > x->key) del(x->r, key); else { if (!(x->l) && !(x->r)) { trans(x, NULL); return; } else if (!(x->l)) trans(x, x->r); else if (!(x->r)) trans(x, x->l); else { node *y = find_min(x->r); del(x->r, y->key); x->key = y->key; } } x->update(); if (x->bf == 2) { if (x->l->bf >= 0) rrot(x); else if (x->l->bf == -1) lrot(x->l), rrot(x); } else if (x->bf == -2) { if (x->r->bf <= 0) lrot(x); else if (x->r->bf == 1) rrot(x->r), lrot(x); } } int main() { ios::sync_with_stdio(false), cin.tie(0), cout.tie(0); root = NULL; int n, x; cin >> n; while (n--) { string s; cin >> s >> x; if (s == "INS") insert(root, new node(x)); else if (s == "DEL") del(root, x); else if (s == "IND") find_kth(root, x); } }
Leave a Comment