variance.py

import numpy as np
import matplotlib.pyplot as plt

input_file = "input_day14.txt"
with open(input_file, "r") as file:
    data = file.readlines()

GRID_WIDTH = 101
GRID_HEIGHT = 103

robots = []
for line in data:
    parts = line.strip().split(" ")
    pos = tuple(map(int, parts[0][2:].strip(",").split(",")))
    vel = tuple(map(int, parts[1][2:].split(",")))
    robots.append((pos, vel))

max_time = 10000
var_x_list = []
var_y_list = []
combined_var_list = []
from tqdm import tqdm
# calculate variances
for t in tqdm(range(max_time + 1)):
    xs = []
    ys = []
    for (px, py), (vx, vy) in robots:
        x = (px + vx * t) % GRID_WIDTH
        y = (py + vy * t) % GRID_HEIGHT
        xs.append(x)
        ys.append(y)
    var_x = np.var(xs)
    var_y = np.var(ys)
    combined_var = var_x + var_y
    var_x_list.append(var_x)
    var_y_list.append(var_y)
    combined_var_list.append(combined_var
min_combined_var = min(combined_var_list)
min_time_step = combined_var_list.index(min_combined_var)
print(min_time_step)