2 files changed, 66 insertions, 35 deletions

diff --git a/day5/__init__.py b/day5/__init__.py
index 4d5260d..b16eba7 100644
--- a/day5/__init__.py
+++ b/day5/__init__.py
@@ -19,27 +19,54 @@ class Assignment(BaseAssignment):
         return list(super().read_input(example))
 
 
-def is_horizontal(v: Vector):
-    return v[0][0] == v[1][0]
-
-def is_vertical(v: Vector):
-    return v[0][1] == v[1][1]
-
-def points_in_vector(vector: Vector) -> List[Tuple[int, int]]:
-    first_x = min(vector[0][0], vector[1][0])
-    last_x = max(vector[0][0], vector[1][0])
-    first_y = min(vector[0][1], vector[1][1])
-    last_y = max(vector[0][1], vector[1][1])
+def is_horizontal(vector: Vector):
+    return vector[0][0] == vector[1][0]
+
+def is_vertical(vector: Vector):
+    return vector[0][1] == vector[1][1]
+
+def is_diagonal(vector: Vector):
+    x_start = vector[0][0]
+    x_end = vector[1][0]
+    y_start = vector[0][1]
+    y_end = vector[1][1]
+
+    return abs(x_end - x_start) == abs(y_end - y_start)
+
+def points_in_vector(vector: Vector, includes_diagonals: bool = False) -> List[Tuple[int, int]]:
+    x_start = vector[0][0]
+    x_end = vector[1][0]
+    y_start = vector[0][1]
+    y_end = vector[1][1]
+    
+    delta_x = x_end - x_start
+    try:
+        direction_x = int(delta_x / abs(delta_x))
+    except ZeroDivisionError:
+        direction_x = 0
+
+    delta_y = y_end - y_start
+    try:
+        direction_y = int(delta_y / abs(delta_y))
+    except ZeroDivisionError:
+        direction_y = 0
 
     return [
-        (vector[0][0], y)
-        for y
-        in range(first_y, last_y + 1)
+        (x_start, y_start + y_delta)
+        for y_delta
+        in range(0, delta_y + direction_y, direction_y)
     ] if is_horizontal(vector) else [
-        (x, vector[0][1])
-        for x
-        in range(first_x, last_x + 1)
-    ] if is_vertical(vector) else []
+        (x_start + x_delta, y_start)
+        for x_delta
+        in range(0, delta_x + direction_x, direction_x)
+    ] if is_vertical(vector) else [
+        (x_start + x_delta, y_start + y_delta)
+        for x_delta, y_delta
+        in zip(
+           range(0, delta_x + direction_x, direction_x),
+           range(0, delta_y + direction_y, direction_y)
+        )
+    ] if includes_diagonals and is_diagonal(vector) else []
 
 
 class AssignmentOne(Assignment):
@@ -54,21 +81,13 @@ class AssignmentOne(Assignment):
         return len([c for c in Counter(coordinates).values() if c >= 2])
 
 
+class AssignmentTwo(Assignment):
+    example_result = 12
 
+    def run(self, input: List[Vector]) -> int:
+        coordinates = []
+
+        for vector in input:
+            coordinates += points_in_vector(vector, True)
 
-# class AssignmentTwo(Assignment):
-#     example_result = 1924
-#
-#     def run(self, input: Tuple[List[int], List[BingoCard]]) -> int:
-#         nrs, cards = input
-#
-#         in_game = copy(cards)
-#         for nr in nrs:
-#             for card in copy(in_game):
-#                 card.mark(nr)
-#
-#                 if card.bingo:
-#                     in_game.remove(card)
-#
-#                 if len(in_game) == 0:
-#                     return nr * sum(card.unmarked)
\ No newline at end of file
+        return len([c for c in Counter(coordinates).values() if c >= 2])
\ No newline at end of file
diff --git a/day5/test_init.py b/day5/test_init.py
index 5096a60..a4d7bff 100644
--- a/day5/test_init.py
+++ b/day5/test_init.py
@@ -7,7 +7,19 @@ def test_points_in_vector():
     ]
 
     assert points_in_vector(((9,7), (7,7))) == [
-        (7, 7), (8, 7), (9, 7)
+        (9, 7), (8, 7), (7, 7)
+    ]
+
+    assert points_in_vector(((1,1), (3,3)), False) == []
+
+    assert points_in_vector(((1,1), (3,3)), True) == [
+        (1, 1), (2, 2), (3, 3)
+    ]
+
+    assert points_in_vector(((9,7), (7,9)), False) == []
+
+    assert points_in_vector(((9,7), (7,9)), True) == [
+        (9, 7), (8, 8), (7, 9),
     ]