acwing算法提高之数据结构--平衡树Treap

1 介绍

本博客用来记录使用平衡树求解的题目。

插入、删除、查询操作的时间复杂度都是O(logN)。

动态维护一个有序序列。

2 训练

题目1：253普通平衡树

C++代码如下，

#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>using namespace std;const int N = 100010, INF = 1e8;int n;
struct Node {int l, r;int key, val;int cnt, size;
}tr[N];int root, idx;void pushup(int p) {tr[p].size = tr[tr[p].l].size + tr[tr[p].r].size + tr[p].cnt;
}int get_node(int key) {tr[++idx].key = key;tr[idx].val = rand();tr[idx].cnt = tr[idx].size = 1;return idx;
}void zig(int &p) {int q = tr[p].l;tr[p].l = tr[q].r, tr[q].r = p, p = q;pushup(tr[p].r), pushup(p);
}void zag(int &p) {int q = tr[p].r;tr[p].r = tr[q].l, tr[q].l = p, p = q;pushup(tr[p].l), pushup(p);
}void build() {get_node(-INF), get_node(INF);root = 1, tr[1].r = 2;pushup(root);if (tr[1].val < tr[2].val) zag(root);
}void insert(int &p, int key) {if (!p) p = get_node(key);else if (tr[p].key == key) tr[p].cnt++;else if (tr[p].key > key) {insert(tr[p].l, key);if (tr[tr[p].l].val > tr[p].val) zig(p);} else {insert(tr[p].r, key);if (tr[tr[p].r].val > tr[p].val) zag(p);}pushup(p);
}void remove(int &p, int key) {if (!p) return;if (tr[p].key == key) {if (tr[p].cnt > 1) tr[p].cnt --;else if (tr[p].l || tr[p].r) {if (!tr[p].r || tr[tr[p].l].val > tr[tr[p].r].val) {zig(p);remove(tr[p].r, key);} else {zag(p);remove(tr[p].l, key);}} else {p = 0;}} else if (tr[p].key > key) remove(tr[p].l, key);else remove(tr[p].r, key);pushup(p);
}int get_rank_by_key(int p, int key) {if (!p) return 0;if (tr[p].key == key) return tr[tr[p].l].size + 1;if (tr[p].key > key) return get_rank_by_key(tr[p].l, key);return tr[tr[p].l].size + tr[p].cnt + get_rank_by_key(tr[p].r, key);
}int get_key_by_rank(int p, int rank) {if (!p) return INF;if (tr[tr[p].l].size >= rank) return get_key_by_rank(tr[p].l, rank);if (tr[tr[p].l].size + tr[p].cnt >= rank) return tr[p].key;return get_key_by_rank(tr[p].r, rank - tr[tr[p].l].size - tr[p].cnt);
}int get_prev(int p, int key) {if (!p) return -INF;if (tr[p].key >= key) return get_prev(tr[p].l, key);return max(tr[p].key, get_prev(tr[p].r, key));
}int get_next(int p, int key) {if (!p) return INF;if (tr[p].key <= key) return get_next(tr[p].r, key);return min(tr[p].key, get_next(tr[p].l, key));
}int main() {build();scanf("%d", &n);while (n--) {int opt, x;scanf("%d%d", &opt, &x);if (opt == 1) insert(root, x);else if (opt == 2) remove(root, x);else if (opt == 3) printf("%d\n", get_rank_by_key(root, x) - 1);else if (opt == 4) printf("%d\n", get_key_by_rank(root, x + 1));else if (opt == 5) printf("%d\n", get_prev(root, x));else printf("%d\n", get_next(root, x));        }return 0;
}

使用C++的multiset，超时了，通过了 6/11个数据

#include <iostream>
#include <cstring>
#include <algorithm>
#include <set>using namespace std;int main() {multiset<int> s;int m;cin >> m;while (m--) {int op, x;cin >> op >> x;if (op == 1) {s.insert(x);} else if (op == 2) {s.extract(x);} else if (op == 3) {int idx = distance(s.begin(), s.lower_bound(x));cout << 1 + idx << endl;} else if (op == 4) {x -= 1; //下标从0开始auto it = s.begin();advance(it, x);cout << *it << endl;} else if (op == 5) {auto it = s.lower_bound(x);it--; //保证一定有解cout << *it << endl;} else if (op == 6) {auto it = s.upper_bound(x);cout << *it << endl;}}return 0;
}

使用python3的库sortedcontainers，发现acwing报错没有该模块，

在这里插入图片描述

from sortedcontainers import SortedList
import bisectsl = SortedList()m = int(input())
for i in range(m):line = intput()ls_line = line.split()ls_line = [int(x) for x in ls_line]op = ls_line[0]x = ls_line[1]if op == 1:sl.add(x)elif op == 2:sl.remove(x)elif op == 3:i = bisect.bisect_left(sl, x)print(i + 1)elif op == 4:print(sl[x-1])elif op == 5:it = bisect.bisect_left(ls, x)it -= 1print(sl[it])elif op == 6:it = biset.bisect_right(ls, x)print(sl[it])

题目2：265营业额统计

使用C++的multiset，

#include <iostream>
#include <cstring>
#include <algorithm>
#include <set>
#include <climits>using namespace std;int main() {int n;cin >> n;long long res = 0;multiset<int> s;for (int i = 0; i < n; ++i) {int t;cin >> t;if (i == 0) {res += t;s.insert(t); //插入数tcontinue;}int ans = 0x3f3f3f3f;//求大于等于t的迭代器auto iter1 = s.lower_bound(t);if (iter1 != s.end()) {int x = *iter1;ans = min(ans, abs(x - t));}if (iter1 != s.begin()) {auto iter2 = iter1;iter2--;int x = *iter2;ans = min(ans, abs(x - t));}res += ans;s.insert(t); //插入数t}cout << res << endl;return 0;
}

C++代码如下，

#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>using namespace std;typedef long long LL;const int N = 33010, INF = 1e7;int n;
struct Node {int l, r;int key, val;
}tr[N];int root, idx;int get_node(int key) {tr[++idx].key = key;tr[idx].val = rand();return idx;
}void build() {get_node(-INF), get_node(INF);root = 1, tr[1].r = 2;
}void zig(int &p) {int q = tr[p].l;tr[p].l = tr[q].r, tr[q].r = p, p = q;
}void zag(int &p) {int q = tr[p].r;tr[p].r = tr[q].l, tr[q].l = p, p = q;
}void insert(int &p, int key) {if (!p) p = get_node(key);else if (tr[p].key == key) return;else if (tr[p].key > key) {insert(tr[p].l, key);if (tr[tr[p].l].val > tr[p].val) zig(p);} else {insert(tr[p].r, key);if (tr[tr[p].r].val > tr[p].val) zag(p);}
}int get_prev(int p, int key) {if (!p) return -INF;if (tr[p].key > key) return get_prev(tr[p].l, key);return max(tr[p].key, get_prev(tr[p].r, key));
}int get_next(int p, int key) {if (!p) return INF;if (tr[p].key < key) return get_next(tr[p].r, key);return min(tr[p].key, get_next(tr[p].l, key));
}int main() {build();scanf("%d", &n);LL res = 0;for (int i = 1; i <= n; ++i) {int x;scanf("%d", &x);if (i == 1) res += x;else res += min(x - get_prev(root, x), get_next(root, x) - x);insert(root, x);}printf("%lld\n", res);return 0;
}