2 files changed, 34 insertions, 4 deletions

diff --git a/day9/__init__.py b/day9/__init__.py
index 4846468..e0f1d61 100644
--- a/day9/__init__.py
+++ b/day9/__init__.py
@@ -59,16 +59,21 @@ class Assignment(BaseAssignment, ABC):
     def next_knot(head: Coordinate, tail: Coordinate):
         delta = head - tail
 
-        if abs(delta.x) > 1 and delta.y == 0:
+        if abs(delta.x) < 2 and abs(delta.y) < 2:
+            pass
+        elif abs(delta.x) > 1 and delta.y == 0:
             tail.x += delta.x - delta.polarity_x
-        if abs(delta.x) > 1 and abs(delta.y) == 1:
+        elif 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:
+        elif 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:
+        elif abs(delta.y) > 1 and abs(delta.x) == 1:
             tail.y += delta.y - delta.polarity_y
             tail.x += delta.x
+        elif abs(delta.x) > 1 and abs(delta.y) > 1:
+            tail.x += delta.x - delta.polarity_x
+            tail.y += delta.y - delta.polarity_y
 
     def tail_positions(self, input: Iterator[str], length: int) -> Iterator[Coordinate]:
         knots = [Coordinate(0, 0) for _ in range(length)]
diff --git a/day9/test_init.py b/day9/test_init.py
index eb889c7..431dbbb 100644
--- a/day9/test_init.py
+++ b/day9/test_init.py
@@ -75,6 +75,31 @@ class TestAssignment:
             assignment.visualize(range(1), range(-5, 1), unique_positions) == expected
         )
 
+    def test_diagonal(self):
+        expected = "\n".join(
+            [
+                "..#",
+                ".#.",
+                "s..",
+            ]
+        )
+
+        input = [
+            "R 1",
+            "U 1",
+            "R 1",
+            "U 1",
+            "R 1",
+            "U 1",
+            "R 1",
+            "U 1",
+        ]
+
+        assignment = day9.AssignmentOne(path=os.path.dirname(day9.__file__))
+        unique_positions = assignment.unique_tail_positions(input=input, length=2)
+
+        assert assignment.visualize(range(3), range(3), unique_positions) == expected
+
 
 class TestAssignmentOne:
     def test_output_visualization(self):