2 files changed, 19 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),
             ]
         )
 
diff --git a/day20/test_init.py b/day20/test_init.py
index b3ccc44..24b94cb 100644
--- a/day20/test_init.py
+++ b/day20/test_init.py
@@ -25,3 +25,11 @@ def test_rotation():
 
     items = AssignmentOne.move(items, 4)
     assert items == [1, 2, -3, 4, 0, 3, -2]
+
+
+def test_get_nth_number_after_0():
+    items = [1, 2, -3, 4, 0, 3, -2]
+
+    assert AssignmentOne.get_nth_number_after_0(items, 1000) == 4
+    assert AssignmentOne.get_nth_number_after_0(items, 2000) == -3
+    assert AssignmentOne.get_nth_number_after_0(items, 3000) == 2