solve Day17 Part One in python

Day17/python/solution1.py

@@ -0,0 +1,109 @@
+from heapq import heappush, heappop
+
+# Constants for movement directions
+DIRECTIONS = [(-1, 0), (1, 0), (0, -1), (0, 1)]
+
+
+def read_grid(filename):
+    """
+    Reads the grid of numbers from the given file.
+    """
+    with open(filename) as f:
+        return [tuple(map(int, line.strip())) for line in f.readlines()]
+
+
+def is_valid_position(row, col, grid):
+    """
+    Checks if the given position is within the bounds of the grid.
+    """
+    return 0 <= row < len(grid) and 0 <= col < len(grid[0])
+
+
+def find_least_heat_loss(grid):
+    """
+    Finds the path with the least heat loss in the given grid.
+    """
+    queue = [(0, 0, 0, 0, 0, 0)]  # heat_loss, row, col, dr, dc, n
+    seen = set()
+
+    while queue:
+        heat_loss, row, col, dr, dc, n = heappop(queue)
+
+        # Check if destination is reached
+        if row == len(grid) - 1 and col == len(grid[0]) - 1:
+            return heat_loss
+
+        if (row, col, dr, dc, n) in seen:
+            continue
+
+        seen.add((row, col, dr, dc, n))
+
+        # Continue in the same direction if not moved more than 3 blocks
+        if n < 3 and (dr, dc) != (0, 0):
+            new_row, new_col = row + dr, col + dc
+            if is_valid_position(new_row, new_col, grid):
+                heappush(
+                    queue,
+                    (
+                        heat_loss + grid[new_row][new_col],
+                        new_row,
+                        new_col,
+                        dr,
+                        dc,
+                        n + 1,
+                    ),
+                )
+
+        # Explore adjacent directions, avoiding immediate reversals
+        for new_dr, new_dc in DIRECTIONS:
+            if (new_dr, new_dc) in ((dr, dc), (-dr, -dc)):
+                continue
+
+            new_row, new_col = row + new_dr, col + new_dc
+            if is_valid_position(new_row, new_col, grid):
+                heappush(
+                    queue,
+                    (
+                        heat_loss + grid[new_row][new_col],
+                        new_row,
+                        new_col,
+                        new_dr,
+                        new_dc,
+                        1,
+                    ),
+                )
+
+    return -1  # Return -1 if no path is found
+
+
+def test_algorithm():
+    """
+    Tests the algorithm with test data and asserts the correct output.
+    """
+    test_grid = read_grid("../test.txt")
+    expected_result = 102  # Expected heat loss
+    result = find_least_heat_loss(test_grid)
+    assert (
+        result == expected_result
+    ), f"Test failed: Expected {expected_result}, got {result}"
+    print("Test passed successfully!")
+
+
+def main():
+    """
+    Main function to run the puzzle solution.
+    """
+    try:
+        # Run test first
+        test_algorithm()
+
+        # If test passes, run the main puzzle
+        grid = read_grid("../input.txt")
+        result = find_least_heat_loss(grid)
+        print(f"Least heat loss for the main puzzle: {result}")
+    except Exception as e:
+        print(f"Error occurred: {e}")
+
+
+if __name__ == "__main__":
+    main()