AOC 2023, day 6

aoc2023_day6.py

import re
from math import floor, ceil, prod
from more_itertools import chunked
from sympy import Symbol, S, solve_univariate_inequality

t = Symbol("t", domain=S.Naturals)
    
def to_range(conditions):
    ineq1, ineq2 = conditions.args
    if ineq1.lhs == t:  # rule: lhs < rhs
        return range(floor(ineq2.lhs) + 1, ceil(ineq1.rhs))
    return range(floor(ineq1.lhs) + 1, ceil(ineq2.rhs))

def compute(n):
    return prod(
        len(to_range(solve_univariate_inequality(-t**2 + t*time > dist, t)))
        for time, dist in zip(*chunked(n, len(n)//2))
    )

part_one = compute(list(map(int, re.findall("\d+", inp))))
part_two = compute(list(map(int, ("".join(re.findall("\d+", s)) for s in inp.splitlines()))))
print(part_one, part_two)