P6136 Enhanced Balanced Tree Template
After struggling with rotational Treaps, I've concluded they're overly complex. FHQ Treaps offer a much cleaner implementation, with roughly half the code size.
Rotational Treap Implementation
Structure and data definitions:
const int INF=1e18;
struct TreeNode {
int left, right;
int value, priority;
int count, size;
} nodes[MAX_NODES];
#define left_child(p) nodes[p].left
#define right_child(p) nodes[p].right
int root, node_count;
Creating a new node and returning its index:
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
inline int build_node(int val) {
nodes[++node_count] = {0, 0, val, (int)rng(), 1, 1};
return node_count;
}
Updating node size:
inline void update_size(int p) {
nodes[p].size = nodes[left_child(p)].size + nodes[right_child(p)].size + nodes[p].count;
return;
}
Right rotation (zig) and left rotation (zag) operations:
inline void right_rotate(int &p) {
int q = left_child(p);
nodes[p].left = right_child(q);
nodes[q].right = p;
p = q;
update_size(right_child(p));
update_size(p);
return;
}
inline void left_rotate(int &p) {
int q = right_child(p);
nodes[p].right = left_child(q);
nodes[q].left = p;
p = q;
update_size(left_child(p));
update_size(p);
return;
}
Building the tree with sentinle nodes:
inline void initialize_tree() {
root = build_node(-INF);
nodes[root].right = build_node(INF);
nodes[right_child(root)].priority = nodes[root].priority - 1;
update_size(root);
return;
}
Getting rank by value:
int find_rank_by_value(int val, int p) {
if (!p) return 1;
if (val == nodes[p].value) return nodes[left_child(p)].size + 1;
if (val < nodes[p].value) return find_rank_by_value(val, left_child(p));
if (val > nodes[p].value) return find_rank_by_value(val, right_child(p)) + nodes[left_child(p)].size + nodes[p].count;
return -INF;
}
Getting value by rank:
int find_value_by_rank(int rank, int p) {
if (!p) return INF;
if (rank <= nodes[left_child(p)].size) return find_value_by_rank(rank, left_child(p));
else if (rank <= nodes[left_child(p)].size + nodes[p].count) return nodes[p].value;
else return find_value_by_rank(rank - nodes[left_child(p)].size - nodes[p].count, right_child(p));
}
Inserting a node:
void insert_node(int val, int &p) {
if (!p) {
p = build_node(val);
return;
}
if (val == nodes[p].value) {
nodes[p].count++;
update_size(p);
return;
}
if (val < nodes[p].value) {
insert_node(val, left_child(p));
if (nodes[p].priority < nodes[left_child(p)].priority) right_rotate(p);
}
if (val > nodes[p].value) {
insert_node(val, right_child(p));
if (nodes[p].priority < nodes[right_child(p)].priority) left_rotate(p);
}
update_size(p);
return;
}
Removing a node:
void remove_node(int val, int &p) {
if (!p) return;
if (val == nodes[p].value) {
if (nodes[p].count > 1) {
nodes[p].count--;
update_size(p);
return;
}
else if (left_child(p) || right_child(p)) {
if (!right_child(p) || nodes[left_child(p)].priority > nodes[right_child(p)].priority)
right_rotate(p), remove_node(val, right_child(p));
else left_rotate(p), remove_node(val, left_child(p));
update_size(p);
}
else p = 0;
return;
}
if (val < nodes[p].value) remove_node(val, left_child(p));
if (val > nodes[p].value) remove_node(val, right_child(p));
update_size(p);
return;
}
Finding predecessor:
int find_predecessor(int val) {
int result = 0, p = root;
while (p) {
if (val == nodes[p].value) {
if (left_child(p) > 0) {
p = left_child(p);
while (right_child(p) > 0) p = right_child(p);
result = p;
}
break;
}
if (nodes[p].value < val && (nodes[p].value > nodes[result].value || !result)) result = p;
p = (val < nodes[p].value) ? left_child(p) : right_child(p);
}
return nodes[result].value;
}
Finding successor:
int find_successor(int val) {
int result = 0, p = root;
while (p) {
if (val == nodes[p].value) {
if (right_child(p)) {
p = right_child(p);
while (left_child(p)) p = left_child(p);
result = p;
}
break;
}
if (nodes[p].value > val && (nodes[p].value < nodes[result].value || !result)) result = p;
p = (val < nodes[p].value) ? left_child(p) : right_child(p);
}
return nodes[result].value;
}
Non-rotational Treap (FHQ Treap)
Structure and data definitions:
struct TreapNode {
int left, right;
int value, priority;
int size;
} nodes[MAX_SIZE];
int node_count, root;
#define left_child(p) nodes[p].left
#define right_child(p) nodes[p].right
Creating a new node:
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
inline int create_node(int val) {
nodes[++node_count] = {0, 0, val, (int)rng(), 1};
return node_count;
}
Updating node size:
inline void update_node_size(int p) {
nodes[p].size = nodes[left_child(p)].size + nodes[right_child(p)].size + 1;
return;
}
Split operation:
void split_tree(int val, int p, int &left_part, int &right_part) {
if (!p) {
left_part = right_part = 0;
return;
}
if (nodes[p].value <= val) {
left_part = p;
split_tree(val, right_child(p), right_child(p), right_part);
update_node_size(left_part);
}
else {
right_part = p;
split_tree(val, left_child(p), left_part, left_child(p));
update_node_size(right_part);
}
return;
}
Merge operation:
int merge_trees(int left_tree, int right_tree) {
if (!left_tree || !right_tree) return left_tree | right_tree;
if (nodes[left_tree].priority > nodes[right_tree].priority) {
right_child(left_tree) = merge_trees(right_child(left_tree), right_tree);
update_node_size(left_tree);
return left_tree;
}
else {
left_child(right_tree) = merge_trees(left_tree, left_child(right_tree));
update_node_size(right_tree);
return right_tree;
}
}
Insert operation:
void insert_value(int val) {
int new_node = create_node(val);
int left_part, right_part;
split_tree(val, root, left_part, right_part);
root = merge_trees(left_part, merge_trees(new_node, right_part));
return;
}
Delete operation:
void delete_value(int val) {
int temp_right, right_part;
split_tree(val, root, temp_right, right_part);
int left_part, middle_part;
split_tree(val - 1, temp_right, left_part, middle_part);
middle_part = merge_trees(left_child(middle_part), right_child(middle_part));
root = merge_trees(left_part, merge_trees(middle_part, right_part));
return;
}
Getting rank by value:
int get_rank_by_value(int val) {
int left_part, right_part;
split_tree(val - 1, root, left_part, right_part);
int result = nodes[left_part].size + 1;
root = merge_trees(left_part, right_part);
return result;
}
Getting value by rank:
int get_value_by_rank(int rank) {
int current = root;
while (current) {
int left_size = nodes[left_child(current)].size;
if (rank == left_size + 1) return nodes[current].value;
if (rank <= left_size) current = left_child(current);
else rank -= left_size + 1, current = right_child(current);
}
return -1;
}
Finding predecessor and successor:
int get_predecessor(int val) {
return get_value_by_rank(get_rank_by_value(val) - 1);
}
int get_successor(int val) {
return get_value_by_rank(get_rank_by_value(val + 1));
}
P3391 Artistic Balanced Tree Template
This problem demonstrates the elegance of FHQ Treaps. In this case, the node values aren't static but represent sequence indices. These indices natural satisfy BST properties: indices to the left are smaller, and those to the right are larger.
Since these indices are dynamic, we don't store them statically but compute them dynamically within the split function. This aproach is highly extensible for interval-based problems.
Key implementation details include lazy propagation for reversal operations:
struct ReversibleNode {
int left, right;
int index, priority;
bool reverse_flag;
int size;
} nodes[MAX_SIZE];
int node_count, root;
inline void propagate_lazy(int p) {
if (nodes[p].reverse_flag) {
nodes[left_child(p)].reverse_flag ^= true;
nodes[right_child(p)].reverse_flag ^= true;
swap(left_child(p), right_child(p));
nodes[p].reverse_flag = false;
}
return;
}
void reverse_interval(int l, int r) {
int temp_right, right_part;
split_tree(r, root, temp_right, right_part);
int left_part, middle_part;
split_tree(l - 1, temp_right, left_part, middle_part);
nodes[middle_part].reverse_flag ^= true;
root = merge_trees(merge_trees(left_part, middle_part), right_part);
return;
}