chore: create new project structure and aoc.py runner script

adventofcode/2024/day14/day14.py

@@ -0,0 +1,106 @@
+from math import prod
+from collections import Counter
+
+
+def parse_bots(lines):
+    lines = [l.rstrip().split(" ") for l in lines]
+    lines = [complex(*map(int, x.split("=")[1].split(","))) for l in lines for x in l]
+    return [lines[i:i+2] for i in range(0, len(lines), 2)] # [(pos, velocity), ...]
+
+
+def simulate_bots(bots, grid_size, steps=100, part2=False):
+    step = 0
+    width, height = grid_size
+    stats = []
+    while step < steps:
+        new_bots = []
+        for pos, velocity in bots:
+            pos = pos + velocity
+            if pos.real >= width:
+                pos -= width
+            if pos.real < 0:
+                pos += width
+            if pos.imag >= height:
+                pos -= height * 1j
+            if pos.imag < 0:
+                pos += height * 1j
+            new_bots.append((pos, velocity))
+        bots = new_bots
+        step += 1
+
+        if part2:
+            # search step which maximizes safety value
+            safety = calculate_safety(Counter([p for p, _ in bots]), grid_size)
+            stats.append((safety, step))
+
+    return [pos for pos, _ in bots], stats
+
+
+def determine_quadrant(pos, grid_size):
+    width, height = grid_size
+    q = None
+    if pos.real < width // 2 and pos.imag < height // 2:
+        q = 0
+    elif pos.real > width // 2 and pos.imag < height //2:
+        q = 1
+    elif pos.real < width // 2 and pos.imag > height // 2:
+        q = 2
+    elif pos.real > width // 2 and pos.imag > height // 2:
+        q = 3
+    return q
+
+
+def calculate_safety(bots, grid_size):
+    total_quadrants = [0, 0, 0, 0]
+    for pos, count in bots.items():
+        q = determine_quadrant(pos, grid_size)
+        if q is None:  # ignore middle row and col
+            continue
+        total_quadrants[q] += count
+    return prod(total_quadrants)
+
+
+def part1(bots, grid_size):
+    bots, _ = simulate_bots(bots, grid_size)
+    c = Counter(bots)
+    return calculate_safety(c, grid_size)
+
+
+def part2(bots, grid_size):
+    max_step = grid_size[0] * grid_size[1]  # input is periodic
+    _, stats = simulate_bots(bots, grid_size, max_step, part2=True)
+    return sorted(stats)[0][1]
+
+
+def main(lines):
+    bots = parse_bots(lines)
+    total = part1(bots, grid_size)
+    print("Part 1: ", total)
+    p2 = part2(bots, grid_size)
+    print("Part 2: ", p2)
+
+
+def print_grid(c, grid_size):
+    width, height = grid_size
+    for y in range(height):
+        for x in range(width):
+            a = c.get(complex(x, y))
+            if a is None:
+                print(".", end="")
+            else:
+                print(a, end="")
+        print()
+
+        
+if __name__ == "__main__":
+    import sys
+    if len(sys.argv) < 2:
+        infile = "example.txt"
+        grid_size = (11, 7)
+    else:
+        infile = sys.argv[1]
+        grid_size = (101, 103)
+    with open(infile) as f:
+        lines = f.readlines()
+        main(lines)
+