add 2024/day14.2

2024/14/main.py

@@ -7,6 +7,7 @@
 # pylint: disable=missing-module-docstring,missing-function-docstring,consider-using-f-string,no-else-return
 
 import re
+import matplotlib.pyplot as plt
 
 
 def readinput():
@@ -25,13 +26,10 @@ def readinput():
 
 
 def get_robot_pos_after_time(area: dict, robot: tuple, seconds: int) -> tuple:
-    for _ in range(seconds):
-        nxtx = (robot["pos"][0] + robot["move"][0]) % area["width"]
-        nxty = (robot["pos"][1] + robot["move"][1]) % area["height"]
+    finx = (robot["pos"][0] + robot["move"][0] * seconds) % area["width"]
+    finy = (robot["pos"][1] + robot["move"][1] * seconds) % area["height"]
 
-        robot["pos"] = nxtx, nxty
-
-    return robot["pos"]
+    return (finx, finy)
 
 
 def get_quadrant(area: dict, position: tuple) -> int:
@@ -63,13 +61,30 @@ def assign_qadrants(area: dict, positions: dict) -> list:
     return quadrant_count
 
 
+def draw_positions_in_area(area: dict, positions: dict, filename: str):
+    grid = []
+    for _ in range(area["height"]):
+        row = [0] * area["width"]
+        grid.append(row)
+
+    for pos in positions.keys():
+        x, y = pos
+        grid[y][x] = 1
+
+    plt.figure()
+    plt.imshow(grid, cmap="gray", interpolation="nearest")
+    plt.axis("off")
+    plt.savefig(f"{filename}.jpg", format="jpg", bbox_inches="tight", pad_inches=0)
+    plt.close()
+
+
 def main():
     robots = readinput()
 
     area = {"width": 101, "height": 103}
-    seconds = 100
 
     # part 1
+    seconds = 100
     positions = {}
     for robot in robots:
         pos = get_robot_pos_after_time(area, robot, seconds)
@@ -82,6 +97,16 @@ def main():
 
     print("Safety factor: %d" % safety_factor)
 
+    # part 2
+    seconds = 10_000
+    for i in range(seconds):
+        positions = {}
+        for robot in robots:
+            pos = get_robot_pos_after_time(area, robot, i)
+            positions[pos] = positions.get(pos, 0) + 1
+        print(f"\rSaving image for position after {i} seconds", end="", flush=True)
+        draw_positions_in_area(area, positions, str("%05d" % i))
+
 
 if __name__ == "__main__":
     main()