Add subtree size updating during treap join

file.hpp

@@ -14,16 +14,17 @@ struct line {
 	string text;
 	line *left, *right;
 
-	line(count_type _p, string _t) : priority(_p), text(_t), size(-1), left(nullptr), right(nullptr) {}
+	line(count_type _p, string _t) : priority(_p), text(_t), size(1), left(nullptr), right(nullptr) {}
 };
 
 // Treap representation of a file
 class file {
 	line* root;
 
-	file() {
+	// Get size uf a subtreap
-		srand(120);
+	int get_size(line* l) {
-		root = nullptr;
+		if(l == nullptr) return 0;
+		return l->size;
 	}
 
 	// Split treap by k-th element
@@ -31,45 +32,56 @@ class file {
 	two_lines split(line *l, count_type k) {
 		if(l == nullptr) return {nullptr, nullptr};
 
-		if(l->size <= k) {
+		if(l->size < k) {
 			auto two = split(l->right, k);
 			l->right = two.first;
 			return {l, two.second};
 		}
 		else {
-			auto two = split(l->left, k - 1 - (l->left == nullptr ? 0 : l->left->size));
+			auto two = split(l->left, k - 1 - get_size(l->left));
 			l->left = two.second;
 			return {two.first, l};
 		}
 	}
 	// Join two treaps
-	// TODO with sizes
+	// Hopefully with working sizes
 	line* join(line *a, line *b) {
 		if(a == nullptr) return b;
 		if(b == nullptr) return a;
 
 		if(a->priority < b->priority) {
+			a->size += b->size;
 			a->right = join(a->right, b);
 			return a;
 		}
 		else {
+			b->size += a->size;
 			b->left = join(a, b->left);
 			return b;
 		}
 	}
 	line* join(two_lines two) { return join(two.first, two.second); }
-
+	// Find k-th line of file
-public:
 	line* find(line* l, count_type k) {
 		if(l == nullptr) return nullptr;
 
-		if(k < l->size)
+		if(k >= get_size(l->left))
 			return find(l->left, k);
-		else if(k > l->size)
+		else if(k > l->size+1)
-			return find(l->right, k - 1 - (l->left == nullptr ? 0 : l->left->size) );
+			return find(l->right, k - 1 - get_size(l->left));
 		else
 			return l;
 	}
+
+public:
+	file() {
+		srand(120);
+		root = nullptr;
+	}
+	// Don't find index k, but k-th line -> +1
+	line* find(count_type k) {return find(root, k+1);}
+
+	// File access
 	void insert(int k, string s) {
 		line *l = new line(rand(), s);
 
@@ -82,6 +94,7 @@ public:
 		join(two.first, l);
 		join(two);
 	}
+	// Line removal
 	void remove(count_type k) {
 		auto two = split(root, k);
 		two.first = split(two.first, k-1).first;