day5.py

import re
from typing import List, Tuple


def get_lines() -> List[str]:
    with open('input.txt', 'r') as f:
        return f.readlines()


def parse_stacks() -> List[List[str]]:
    lines = get_lines()
    n_stacks = len(lines[0]) // 4
    stacks = [[] for _ in range(n_stacks)]
    for line in lines:
        if '[' not in line or line == '\n':
            break

        for i in range(0, len(line), 4):
            s = line[i:i+4]
            if len(s) == 4 and s[1] != ' ':
                stacks[i//4].append(s[1])

    for stack in stacks:
        stack.reverse()

    return stacks


def parse_moves() -> List[Tuple[int, int]]:
    moves = []
    lines = get_lines()
    pattern = re.compile(r'\d+')
    for line in lines:
        if not line.startswith('move'):
            continue

        n, src, dst = map(int, pattern.findall(line))

        # input is 1-indexed
        moves.append((n, src - 1, dst - 1))

    return moves


def move_part1(stacks, n, src, dst) -> None:
    for _ in range(n):
        stacks[dst].append(stacks[src].pop())


def move_part2(stacks, n, src, dst) -> None:
    a, b = stacks[src][:-n], stacks[src][-n:]
    stacks[src] = a
    stacks[dst] += b
        

def part1():
    stacks = parse_stacks()
    moves = parse_moves()
    for n, src, dst in moves:
        move_part1(stacks, n, src, dst)

    return ''.join([stack[-1] for stack in stacks])


def part2():
    stacks = parse_stacks()
    moves = parse_moves()
    for n, src, dst in moves:
        move_part2(stacks, n, src, dst)

    return ''.join([stack[-1] for stack in stacks])