一个 treap 实现,出现段错误。insert
方法中不调用rotate
是正常(即此时为朴素BST)。
#include <bits/stdc++.h>
struct Node;
typedef Node *NodePtr;
struct Node {
int key, value;
NodePtr parent, children[2];
size_t size;
int priority;
enum Direction {
LEFT = 0,
RIGHT = 1
};
Node(int key)
: key(key), value(1), parent(nullptr),
children{nullptr, nullptr}, size(1),
priority(rand()) {}
void upsize();
void rotate(const Direction);
void insert(const int);
int get_key(const size_t);
};
inline void Node::upsize() {
this->size = this->value;
if (children[LEFT] != nullptr)
this->size += children[LEFT]->size;
if (children[RIGHT] != nullptr)
this->size += children[RIGHT]->size;
}
inline void Node::rotate(const Direction dir) {
/* e.g.
* P P
* | |
* A C
* / \ left / \
* B C ----> A E
* / \ / \
* D E B D
*/
const auto A = this, P = this->parent,
C = this->children[!dir],
D = this->children[!dir]->children[dir];
if (P != nullptr) P->children[A->value > P->value] = C;
A->children[!dir] = D;
C->children[dir] = A;
}
inline void Node::insert(const int key) {
if (this->key == key) {
this->key += 1;
this->upsize();
return;
}
Direction dir = key > this->key ? RIGHT : LEFT;
auto &child = this->children[dir];
if (child != nullptr) {
child->insert(key);
} else {
child = new Node(key);
}
if (this->children[dir]->priority < this->priority) {
this->rotate(!dir ? RIGHT : LEFT);
this->upsize();
}
child->upsize(), this->upsize();
}
inline int Node::get_key(const size_t rank) {
auto &l_child = this->children[LEFT],
&r_child = this->children[RIGHT];
bool pre = l_child != nullptr;
if (pre && rank <= l_child->size)
return l_child->get_key(rank);
else if ((pre && rank <= l_child->size + this->value) ||
rank <= (unsigned)this->value)
return this->key;
else
return r_child->get_key(
rank - (pre ? l_child->size : 0) - this->value);
}
int main() {
int n, t;
std::cin >> n;
std::cin >> t;
Node a = {t};
for (int i = 1; i < n; i++) {
std::cin >> t;
a.insert(t);
}
std::cin >> t;
std::cout << a.get_key(t);
return 0;
}