advent-of-code/2024/day14/day14.py
2025-01-04 14:13:48 +01:00

107 lines
2.8 KiB
Python

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)