1 files changed, 28 insertions, 21 deletions

diff --git a/day9/__init__.py b/day9/__init__.py
index f198deb..4846468 100644
--- a/day9/__init__.py
+++ b/day9/__init__.py
@@ -34,6 +34,8 @@ class Coordinate:
 
 
 class Assignment(BaseAssignment, ABC):
+    rope_length = NotImplemented
+
     def parse_line(self, item: str) -> Tuple[str, int]:
         direction, amount = item.split(" ")
         return direction, int(amount)
@@ -54,7 +56,7 @@ class Assignment(BaseAssignment, ABC):
                 head.y -= 1
 
     @staticmethod
-    def next_tail(head: Coordinate, tail: Coordinate):
+    def next_knot(head: Coordinate, tail: Coordinate):
         delta = head - tail
 
         if abs(delta.x) > 1 and delta.y == 0:
@@ -62,37 +64,43 @@ class Assignment(BaseAssignment, ABC):
         if abs(delta.x) > 1 and abs(delta.y) == 1:
             tail.x += delta.x - delta.polarity_x
             tail.y += delta.y
-        if abs(delta.y > 1) and delta.x == 0:
+        if abs(delta.y) > 1 and delta.x == 0:
             tail.y += delta.y - delta.polarity_y
         if abs(delta.y) > 1 and abs(delta.x) == 1:
             tail.y += delta.y - delta.polarity_y
             tail.x += delta.x
 
-    def tail_positions(self, input: Iterator[str]) -> Iterator[Coordinate]:
-        head = Coordinate(0, 0)
-        tail = Coordinate(0, 0)
+    def tail_positions(self, input: Iterator[str], length: int) -> Iterator[Coordinate]:
+        knots = [Coordinate(0, 0) for _ in range(length)]
 
         for line in input:
             direction, amount = self.parse_line(line)
 
             for _ in range(amount):
-                self.next_head(head, direction)
-                self.next_tail(head, tail)
-                yield tail
+                for index in range(length):
+                    if index == 0:
+                        self.next_head(knots[index], direction)
+                        continue
+
+                    self.next_knot(knots[index - 1], knots[index])
+
+                yield knots[length - 1]
 
-    def unique_tail_positions(self, input: Iterator[str]) -> Set[Coordinate]:
+    def unique_tail_positions(
+        self, input: Iterator[str], length: int
+    ) -> Set[Coordinate]:
         unique_tail_positions = set()
 
-        for position in self.tail_positions(input):
+        for position in self.tail_positions(input, length):
             unique_tail_positions.add(copy(position))
 
         return unique_tail_positions
 
-    def visualize(self, width: int, height: int, positions: Set[Coordinate]):
+    def visualize(self, width: range, height: range, positions: Set[Coordinate]):
         rows = []
-        for y in range(height):
+        for y in height:
             items = []
-            for x in range(width):
+            for x in width:
                 if x == 0 and y == 0:
                     items.append("s")
                 elif Coordinate(x, y) in positions:
@@ -103,17 +111,16 @@ class Assignment(BaseAssignment, ABC):
 
         return "\n".join(["".join(row) for row in reversed(rows)])
 
+    def run(self, input: Iterator) -> int:
+        unique_positions = self.unique_tail_positions(input, length=self.rope_length)
+        return len(unique_positions)
+
 
 class AssignmentOne(Assignment):
     example_result = 13
-
-    def run(self, input: Iterator) -> int:
-        unique_positions = self.unique_tail_positions(input)
-        print()
-        print(self.visualize(6, 5, unique_positions))
-
-        return len(unique_positions)
+    rope_length = 2
 
 
 class AssignmentTwo(Assignment):
-    pass
+    example_result = 36
+    rope_length = 10