Update the RLE-Encoded problem

notdone/rle-encoded.py

@@ -1,15 +1,19 @@
+# https://coderun.yandex.ru/problem/find-rle-string-length
+
+import unittest
+from abc import ABC, abstractmethod
 from itertools import count
 
 
 def repr_length(number):
     return 0 if number == 1 else len(str(number))
 
 
-class IntervalNode:
-    def __init__(self, left, count):
+class IntervalNode(ABC):
+    def __init__(self, left, length):
         self.left = left
-        self.count = count
-        self.right = left + count - 1
+        self.length = length
+        self.right = left + length - 1
         self.repr_length = 0
 
     @staticmethod
@@ -30,6 +34,7 @@ def gen():
 
         return IntervalBranch(left, list(gen()))
 
+    @abstractmethod
     def find_leaf(self, target):
         assert self.left <= target <= self.right
 
@@ -56,28 +61,18 @@ def fun(*p):
 
 
 class IntervalLeaf(IntervalNode):
-    def __init__(self, left, ch, count):
-        super().__init__(left, count)
-        self.repr_length = repr_length(count) + 1
+    def __init__(self, left, ch, length):
+        super().__init__(left, length)
+        self.repr_length = repr_length(length) + 1
 
     def find_leaf(self, target):
         super().find_leaf(target)
         return self
 
-    #    @debug
-    def len_between(self, l, r):
-        if l <= self.left <= self.right <= r:
-            return self.repr_length
-        l = max(l, self.left)
-        r = min(r, self.right)
-        if r >= l:
-            return 1 + repr(r - l + 1)
-        return 0
-
 
 class IntervalBranch(IntervalNode):
     def __init__(self, left, nodes):
-        super().__init__(left, sum(node.count for node in nodes))
+        super().__init__(left, sum(node.length for node in nodes))
         self.nodes = nodes
         self.repr_length = sum(node.repr_length for node in nodes)
 
@@ -87,10 +82,10 @@ def find_leaf(self, target):
         return node.find_leaf(target)
 
 
-def main():
+def rle_encoded(string, queries):
     def gen_pairs():
         n = 0
-        for c in input():
+        for c in string:
             if c.isdigit():
                 n = 10 * n + int(c)
             else:
@@ -100,11 +95,62 @@ def gen_pairs():
 
     pairs = list(gen_pairs())
     tree = IntervalNode.from_pairs(1, pairs)
+    return [tree.len_between(*query) for query in queries]
 
-    q = int(input())
-    for _ in range(q):
-        print(tree.len_between(*map(int, input().split())))
+
+class TestHelpers(unittest.TestCase):
+    def test_repr_length(self):
+        self.assertEqual(repr_length(1), 0)
+        self.assertEqual(repr_length(10), 2)
+        self.assertEqual(repr_length(100), 3)
+        self.assertEqual(repr_length(123456), 6)
+
+
+# Test the class `IntervalLeaf`
+class TestIntervalMethods(unittest.TestCase):
+    def test_leaf_find_leaf(self):
+        leaf = IntervalLeaf(0, "a", 5)
+        self.assertEqual(leaf.find_leaf(3), leaf)
+
+    def test_branch_find_leaf(self):
+        leaf1 = IntervalLeaf(0, "a", 5)
+        leaf2 = IntervalLeaf(5, "b", 5)
+        branch = IntervalBranch(0, [leaf1, leaf2])
+
+        self.assertEqual(branch.find_leaf(2), leaf1)
+        self.assertEqual(branch.find_leaf(5), leaf2)
+        self.assertEqual(branch.find_leaf(7), leaf2)
+
+    def test_len_between(self):
+        node1 = IntervalLeaf(0, "a", 5)
+        node2 = IntervalLeaf(5, "b", 5)
+        branch = IntervalBranch(0, [node1, node2])
+
+        self.assertEqual(branch.len_between(0, 9), 4)
+
+
+class TestRleEncoded(unittest.TestCase):
+    def test_rle_encoded_1(self):
+        s = "10a"
+        qs = [[1, 1], [3, 7], [1, 10]]
+        result = rle_encoded(s, qs)
+        self.assertListEqual(result, [1, 2, 3])
+
+    def test_rle_encoded_2(self):
+        s = "3a2b"
+        qs = [[1, 1], [3, 4], [1, 5]]
+        result = rle_encoded(s, qs)
+        self.assertListEqual(result, [1, 2, 4])
+
+    def test_rle_encoded_3(self):
+        s = "3ab2c"
+        qs = [[1, 3], [1, 4], [4, 5], [5, 5]]
+        result = rle_encoded(s, qs)
+        self.assertListEqual(result, [2, 3, 2, 1])
 
 
 if __name__ == "__main__":
-    main()
+    s = input()
+    q = int(input())
+    qs = [[int(v) for v in input().split()] for _ in range(q)]
+    print(*rle_encoded(s, qs), sep="\n")