mirror of
https://github.com/thib8956/advent-of-code.git
synced 2025-01-15 14:21:06 +00:00
107 lines
2.8 KiB
Python
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)
|
|
|