From 02bdb353a4a2ca52946b86da45f03e5009df08c9 Mon Sep 17 00:00:00 2001
From: Matuush <matus1@pull.cz>
Date: Fri, 28 Mar 2025 16:16:47 +0100
Subject: [PATCH] Add subtree size updating during treap split

---
 file.hpp | 26 ++++++++++++++++----------
 1 file changed, 16 insertions(+), 10 deletions(-)

diff --git a/file.hpp b/file.hpp
index d6058f7..b1bd35b 100644
--- a/file.hpp
+++ b/file.hpp
@@ -28,19 +28,21 @@ class file {
 	}
 
 	// Split treap by k-th element
-	// TODO with sizes
+	// Hopefully with sizes
 	two_lines split(line *l, count_type k) {
 		if(l == nullptr) return {nullptr, nullptr};
 
-		if(l->size < k) {
-			auto two = split(l->right, k);
-			l->right = two.first;
-			return {l, two.second};
+		if(get_size(l->left) > k) {
+			// In the left subtree
+			auto two = split(l->left, k);
+			l->size -= two.first->size;
+			return {two.first, l};
 		}
 		else {
-			auto two = split(l->left, k - 1 - get_size(l->left));
-			l->left = two.second;
-			return {two.first, l};
+			// In the right subtree
+			auto two = split(l->right, k - (1+get_size(l->left)));
+			l->size -= two.second->size;
+			return {l, two.second};
 		}
 	}
 	// Join two treaps
@@ -68,7 +70,7 @@ class file {
 		if(k >= get_size(l->left))
 			return find(l->left, k);
 		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
 			return l;
 	}
@@ -79,7 +81,11 @@ public:
 		root = nullptr;
 	}
 	// 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
 	void insert(int k, string s) {