RLE Iterator (solved)

Azureki · Sep 9, 2018 · c5435db · c5435db
1 parent aa95a10
commit c5435db
Showing 1 changed file with 25 additions and 13 deletions.
diff --git a/lc900.py b/lc900.py
@@ -8,25 +8,37 @@ def __init__(self, A):
         """
         :type A: List[int]
         """
-        self.index = -1
-        self.A = []
-        for i in range(len(A))[::2]:
-            self.A += [A[i + 1]] * A[i]
+        self.index = 0
+        self.A = A
+        self.len_a = len(A)
 
     def next(self, n):
         """
         :type n: int
         :rtype: int
         """
-        self.index += n
-        if self.index > len(self.A) - 1:
-            return - 1
-        return self.A[self.index]
+        for i in range(self.index, self.len_a, 2):
+            mask = n - self.A[i]
+            if mask > 0:
+                n = mask
+                self.index = i + 2
+                continue
+            elif mask < 0:
+                self.A[i] -= n
+                self.index = i
+                return self.A[i + 1]
+            else:
+                self.index = i + 2
+                return self.A[i + 1]
+        else:
+            return -1
 
 
-A = [3, 8, 0, 9, 2, 5]
+A = [635, 606, 576, 391, 370, 981, 36, 21, 961, 185, 124, 210, 801, 937, 22, 426, 101, 260, 221, 647, 350, 180, 504, 39, 101,
+     989, 199, 358, 732, 839, 919, 169, 673, 967, 58, 676, 846, 342, 250, 597, 442, 174, 472, 410, 569, 509, 311, 357, 838, 251]
 s = RLEIterator(A)
-print(s.next(2))
-print(s.next(1))
-print(s.next(1))
-print(s.next(2))
+
+B = [[5039], [62], [3640], [671], [67], [395], [262], [839], [74], [43], [42], [77], [13], [6], [3], [1], [1], [1], [1], [1], [1], [1], [
+    1], [1], [1], [1], [1], [1], [1], [1], [1], [1], [1], [1], [1], [1], [1], [1], [1], [1], [1], [1], [1], [1], [1], [1], [1], [1], [1], [1]]
+for b in B:
+    print(s.next(b[0]),end=',')