Day 20 [WIP]

1 files changed, 11 insertions, 10 deletions

diff --git a/day20/__init__.py b/day20/__init__.py
index 02ff80d..a61253f 100644
--- a/day20/__init__.py
+++ b/day20/__init__.py
@@ -30,29 +30,30 @@ class AssignmentOne(Assignment):
         if item < 0 and index + item <= 0:
             new_index -= 1
 
+        if new_index == -1:
+            new_index = len(items)
+
         items.insert(new_index, item)
 
         return items
 
+    @staticmethod
+    def get_nth_number_after_0(items: List[int], n: int):
+        index_of_0 = items.index(0)
+        return items[(index_of_0 + n) % len(items)]
+
     def run(self, input: Iterator[I]) -> T:
         input_list = list(input)
-        input_size = len(input_list)
         working_list = list(input_list)
 
         for index, item in enumerate(input_list):
             working_list = self.move(working_list, item)
 
-        print(
-            working_list[1000 % input_size],
-            working_list[2000 % input_size],
-            working_list[3000 % input_size],
-        )
-
         return sum(
             [
-                working_list[1000 % input_size],
-                working_list[2000 % input_size],
-                working_list[3000 % input_size],
+                self.get_nth_number_after_0(working_list, 1000),
+                self.get_nth_number_after_0(working_list, 2000),
+                self.get_nth_number_after_0(working_list, 3000),
             ]
         )