Advent of Code 2024 day 14 part 2

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 expand_blob(todo):
    '''The size of one contiguous group of pixels in todo.'''
    frontier = [todo.pop()]
    size = 1
    while len(frontier) > 0:
        x, y = frontier.pop()
        for dx, dy in ((0, -1), (1, 0), (0, 1), (-1, 0)):
            neighbour = (x + dx, y + dy)
            if neighbour in todo:
                frontier.append(neighbour)
                todo.remove(neighbour)
                size += 1

    return size


def find_blob(coords):
    '''The size of the largest contiguous group of pixels in coords.'''
    todo = set(coords)
    largest = 0
    while len(todo) > 0:
        largest = max(largest, expand_blob(todo))

    # if the blob is larger than the threshold, it might be a christmas tree...
    return largest >= 50


def write_image(coords, width, height, time):
    '''Write the grid to a file as a ppm image, points in coords will be white
       and the rest black.'''
    with open(f'{time}.ppm', 'w') as f:
        print(f'P3\n{width} {height}\n255', file=f)

        for y in range(height):
            for x in range(width):
                val = '255' if (x, y) in coords else '0'
                print(val, val, val, file=f)


def main(robots, width, height, time, max_candidates):
    '''find potential christmas trees, write them to disk as images, and return
       the time they appear.'''
    candidates = list()
    for gen in range(time):
        if candidates == max_candidates:
            break

        coords = {position(*start, *velocity, width, height, gen)
                  for start, velocity in robots}

        # a christmas tree is basically a big blob
        if find_blob(coords):
            # if candidate, write to file as ppm image
            write_image(coords, width, height, gen)
            candidates.append(gen)

    return candidates


parser = ArgumentParser()
parser.add_argument('src', type=FileType('r'))
parser.add_argument('max_candidates', nargs='?', type=int, default=50)
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=10000)

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