Solve puzzle 2 of day 5

Author: bielsnohr

05/hydrothermal_venture.py

@@ -9,59 +9,60 @@ class VentMap():
     def __init__(self, size: int = 9):
         self.map = np.zeros((size, size), dtype=np.int64)
 
-    def mark_line(self, endpoints: npt.ArrayLike) -> None:
+    def mark_line(self, endpoints: npt.ArrayLike, diagonal: bool = False) -> None:
         """Mark the points that a horizontal or vertical line occupies on the map
 
         Args:
             endpoints (npt.ArrayLike): A list of the end point coordinates: [x1, y1, x2, y2]
         """
+        endpoints = np.asarray(endpoints)
         if endpoints[0] == endpoints[2] or endpoints[1] == endpoints[3]:
-            endpoints = np.asarray(endpoints)
             x1, x2 = np.sort(endpoints[[0, 2]])
             y1, y2 = np.sort(endpoints[[1, 3]])
             self.map[y1:y2+1, x1:x2+1] += 1
+        elif diagonal:
+            x1, x2 = endpoints[[0, 2]]
+            y1, y2 = endpoints[[1, 3]]
+            dx = np.abs(x2 - x1)
+            dy = np.abs(y2 - y1)
+            if dx == dy:
+                xrange = np.linspace(start=x1, stop=x2, num=dx+1, dtype=np.int64)
+                yrange = np.linspace(start=y1, stop=y2, num=dy+1, dtype=np.int64)
+                self.map[yrange, xrange] += 1
 
-    def mark_all_lines(self, endpoints_list: npt.ArrayLike) -> None:
+    def mark_all_lines(self, endpoints_list: npt.ArrayLike, diagonal: bool = False) -> None:
         """Mark the points that a horizontal or vertical line occupies on the map
 
         Args:
             endpoints (npt.ArrayLike): A list of lists of the end point coordinates: [[x1, y1, x2, y2], [x1, y1, x2, y2], ...]
         """
         for line in endpoints_list:
-            self.mark_line(endpoints=line)
+            self.mark_line(endpoints=line, diagonal=diagonal)
 
     def calculate_overlap(self):
         return (self.map >= 2).sum()
 
 
-def puzzle_1(inputfile: Path) -> int:
+def puzzle_1(inputfile: Path, diagonal=False) -> int:
     line_parse = re.compile(r'(\d+),(\d+)\s+->\s+(\d+),(\d+)')
     endpoints_list = []
     with open(inputfile, 'r') as file:
         for line in file:
             endpoints_list.append([int(x) for x in line_parse.search(line).groups()])
     endpoints_list = np.array(endpoints_list)
     vent_map = VentMap(size=endpoints_list.max() + 1)
-    vent_map.mark_all_lines(endpoints_list)
+    vent_map.mark_all_lines(endpoints_list, diagonal=diagonal)
     return vent_map.calculate_overlap()
 
 
 def puzzle_2(inputfile: Path) -> int:
-    pass
-    # with open(inputfile, 'r') as file:
-    #     draw = [int(x) for x in file.readline().strip('\n').split(',')]
-    #     file.readline()
-    #     board_strings = []
-    #     for line in file:
-    #         board_strings.append(line.strip('\n'))
-    # bingo_boards = BingoBoardCollection(boards=board_strings)
-    # return bingo_boards.find_losing_score(draw=draw)
+    return puzzle_1(inputfile=inputfile, diagonal=True)
 
 
 def main():
     inputfile = Path("puzzle_1_input.txt")
     print(f"Puzzle 1 Answer = {puzzle_1(inputfile)}")
-    # print(f"Puzzle 2 Answer = {puzzle_2(inputfile)}")
+    print(f"Puzzle 2 Answer = {puzzle_2(inputfile)}")
 
 
 if __name__ == "__main__":

05/test_hydrothermal_venture.py

@@ -1,3 +1,4 @@
+from numpy.core.fromnumeric import diagonal
 from hydrothermal_venture import VentMap, np, Path
 from hydrothermal_venture import puzzle_1, puzzle_2
 import pytest
@@ -60,6 +61,35 @@ def endstate_map():
                      [2, 2, 2, 1, 1, 1, 0, 0, 0, 0]])
 
 
+@pytest.fixture
+def diagonal_mark():
+    return ([5, 5, 8, 2],
+            np.array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                      [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                      [0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
+                      [0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
+                      [0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
+                      [0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
+                      [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                      [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                      [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                      [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]))
+
+
+@pytest.fixture
+def endstate_diagonal_map():
+    return np.array([[1, 0, 1, 0, 0, 0, 0, 1, 1, 0],
+                     [0, 1, 1, 1, 0, 0, 0, 2, 0, 0],
+                     [0, 0, 2, 0, 1, 0, 1, 1, 1, 0],
+                     [0, 0, 0, 1, 0, 2, 0, 2, 0, 0],
+                     [0, 1, 1, 2, 3, 1, 3, 2, 1, 1],
+                     [0, 0, 0, 1, 0, 2, 0, 0, 0, 0],
+                     [0, 0, 1, 0, 0, 0, 1, 0, 0, 0],
+                     [0, 1, 0, 0, 0, 0, 0, 1, 0, 0],
+                     [1, 0, 0, 0, 0, 0, 0, 0, 1, 0],
+                     [2, 2, 2, 1, 1, 1, 0, 0, 0, 0]])
+
+
 @pytest.mark.parametrize(("line, line_mark"), [
     ([0, 9, 5, 9], line2_mark), ([2, 2, 2, 1], line1_mark)
 ])
@@ -80,7 +110,22 @@ def test_ventmap_calculate_overlap(vent_map, endstate_map):
 
 def test_puzzle_1():
     assert puzzle_1(Path(__file__).parent / "puzzle_1_test_input.txt") == 5
-# 
-# 
+
+
+def test_ventmap_diagonal_mark(vent_map, diagonal_mark):
+    endpoints, endstate = diagonal_mark
+    vent_map.mark_line(endpoints=endpoints, diagonal=True)
+    assert np.equal(vent_map.map, endstate).all()
+    
+
+def test_ventmap_diagonal_mark_all_lines(vent_map, list_of_lines, endstate_diagonal_map):
+    vent_map.mark_all_lines(endpoints_list=list_of_lines, diagonal=True)
+    assert np.equal(vent_map.map, endstate_diagonal_map).all()
+
+
+def test_puzzle_2():
+    assert puzzle_2(Path(__file__).parent / "puzzle_1_test_input.txt") == 12
+
+#
 # def test_puzzle_2():
 #     assert puzzle_2(Path(__file__).parent / "puzzle_1_test_input.txt") == 1924