Advent of Code 2024 day 14 part 1

from math import floor, ceil, prod
from argparse import ArgumentParser, FileType


def parse_number(line, index):
    '''Parse the (potentially negative) number starting at line[index] and the
       first index after it.'''
    if line[index] == '-':
        acc = '-'
        index += 1
    else:
        acc = ''

    while index < len(line) and line[index].isdigit():
        acc += line[index]
        index += 1

    return int(acc), index


def parse_pair(line, index):
    '''Parse a pair of numbers separated by a single character and the first
       index after them.'''
    x, index = parse_number(line, index)
    y, index = parse_number(line, index + 1)  # skip the comma

    return (x, y), index


def parse_line(line):
    '''Parse a line of input into the start position and velocity of a
       robot.'''
    # skip past 'p=' and parse the start coords
    start, index = parse_pair(line, 2)

    # skip past ' v=' and parse the velocity
    velocity, _ = parse_pair(line, index + 3)

    return start, velocity


def parse(src):
    '''Parse the input into a list of robots.'''
    return (parse_line(line) for line in src.split('\n') if line != '')


def position(x0, y0, dx, dy, width, height, time):
    '''Return the position of a robot in the grid at the given time.'''
    return (x0 + dx * time) % width, (y0 + dy * time) % height


def inside(coords, min_coords, max_coords):
    '''Test if the coords are in range.'''
    x, y = coords
    min_x, min_y = min_coords
    max_x, max_y = max_coords

    return min_x <= x < max_x and min_y <= y < max_y


def main(robots, width, height, time):
    '''Calculate the product of the populations of each quadrant after time
       steps.'''
    middle_x = width / 2
    middle_y = height / 2

    boundaries = [
        ((0, 0), (floor(middle_x), floor(middle_y))),  # top left
        ((ceil(middle_x), 0), (width, floor(middle_y))),  # top right
        ((0, ceil(middle_y)), (floor(middle_x), height)),  # bottom left
        ((ceil(middle_x), ceil(middle_y)), (width, height))  # bottom right
    ]

    totals = [0, 0, 0, 0]
    for start, velocity in robots:
        coords = position(*start, *velocity, width, height, time)
        for n, (min_coords, max_coords) in enumerate(boundaries):
            if inside(coords, min_coords, max_coords):
                totals[n] += 1

    return prod(totals)


parser = ArgumentParser()
parser.add_argument('src', type=FileType('r'))
parser.add_argument('width', nargs='?', type=int, default=101)
parser.add_argument('height', nargs='?', type=int, default=103)
parser.add_argument('time', nargs='?', type=int, default=100)

if __name__ == '__main__':
    args = parser.parse_args()
    robots = parse(args.src.read())
    print(main(robots, args.width, args.height, args.time))