AOC2021 Day 13 (Python)

day13_tests.py

import re


def read_line(line):
    if m := re.match(r'(\d+),(\d+)', line):
        x_str, y_str = m.groups()
        return (int(x_str), int(y_str)), None
    elif m := re.match(r'fold along ([xy])=(\d+)', line):
        axis, pos = m.groups()
        return None, (axis, int(pos))
    else:
        return None, None


def read_file(lines):
    parsed = [read_line(line.strip()) for line in lines]
    points = set(coord for coord, fold in parsed if coord)
    folds = [fold for coord, fold in parsed if fold]
    return points, folds


def reflect(coord, pos):
    return coord if coord <= pos else pos - (coord - pos)


def do_fold(points, fold):
    axis, pos = fold
    if axis == 'x':
        return set((reflect(x, pos), y) for x, y in points)
    else:
        return set((x, reflect(y, pos)) for x, y in points)


def test_first_fold():
    with open("puzzle_inputs/day13.txt") as file:
        points, folds = read_file(file)

    print()
    print(len(do_fold(points, folds[0])))


def test_all_folds():
    with open("puzzle_inputs/day13.txt") as file:
        points, folds = read_file(file)

    for fold in folds:
        points = do_fold(points, fold)

    print()
    for y in range(max(y for x, y in points) + 1):
        s = ''
        for x in range(max(x for x, y in points) + 1):
            if (x, y) in points:
                s += '#'
            else:
                s += ' '
        print(s)