Add subtree size updating during treap split

This commit is contained in:
Matúš Púll 2025-03-28 16:16:47 +01:00
parent 7f528796c0
commit 02bdb353a4

View file

@ -28,19 +28,21 @@ class file {
} }
// Split treap by k-th element // Split treap by k-th element
// TODO with sizes // Hopefully with sizes
two_lines split(line *l, count_type k) { two_lines split(line *l, count_type k) {
if(l == nullptr) return {nullptr, nullptr}; if(l == nullptr) return {nullptr, nullptr};
if(l->size < k) { if(get_size(l->left) > k) {
auto two = split(l->right, k); // In the left subtree
l->right = two.first; auto two = split(l->left, k);
return {l, two.second}; l->size -= two.first->size;
return {two.first, l};
} }
else { else {
auto two = split(l->left, k - 1 - get_size(l->left)); // In the right subtree
l->left = two.second; auto two = split(l->right, k - (1+get_size(l->left)));
return {two.first, l}; l->size -= two.second->size;
return {l, two.second};
} }
} }
// Join two treaps // Join two treaps
@ -68,7 +70,7 @@ class file {
if(k >= get_size(l->left)) if(k >= get_size(l->left))
return find(l->left, k); return find(l->left, k);
else if(k > l->size+1) else if(k > l->size+1)
return find(l->right, k - 1 - get_size(l->left)); return find(l->right, k - (1+get_size(l->left)) );
else else
return l; return l;
} }
@ -79,7 +81,11 @@ public:
root = nullptr; root = nullptr;
} }
// Don't find index k, but k-th line -> +1 // Don't find index k, but k-th line -> +1
line* find(count_type k) {return find(root, k+1);} line* find(count_type k) {
if(k >= root->size)
return nullptr;
return find(root, k+1);
}
// File access // File access
void insert(int k, string s) { void insert(int k, string s) {