AOC 2021 Day 5

# aoc202105.py

import pathlib
import re
import sys
import numpy as np
from dataclasses import dataclass


@dataclass
class Point:
    x: int
    y: int


@dataclass
class Vent:
    start: Point
    end: Point


def intify(list_of_strings: list[str]) -> list[int]:
    """ Convert a list of strings into one of ints """
    return [int(c) for c in list_of_strings]


def arrange_coords(x1: int, y1: int, x2: int, y2: int) -> tuple[int]:
    """ Make sure coordinates go from L -> R and U -> D """
    if x1 > x2:
        x1, y1, x2, y2 = x2, y2, x1, y1
    elif x1 == x2 and y1 > y2:
        y1, y2 = y2, y1
    return x1, y1, x2, y2


def parse(puzzle_input: str) -> list[int]:
    """ Parse input """
    data = []
    lines = puzzle_input.splitlines()
    pattern = re.compile(r'(\d+),(\d+) -> (\d+),(\d+)')
    for line in lines:
        match = re.search(pattern, line)
        nums = intify(match.groups())
        x1, y1, x2, y2 = arrange_coords(*nums)
        start, end = Point(x1, y1), Point(x2, y2)
        data.append(Vent(start, end))
    return data


def get_straights(lines: list['Vent']) -> list['Vent']:
    """ Return a list that meets the critieria in spec:
        namely start.x == end.x or start.y == end.y """
    straights = []
    for line in lines:
        if any([line.start.x == line.end.x, line.start.y == line.end.y]):
            straights.append(line)
    return straights


def is_diagonal(x1: int, y1: int, x2: int, y2: int) -> bool:
    """ Return true if line is diagonal """
    return any([x1 + y2 == y1 + x2, x1 + x2 == y1 + y2, x1 + y1 == x2 + y2])


def get_diagonals(lines: list['Vent']) -> list['Vent']:
    """ get 45 degree diagonal vents """
    diagonals = []
    for line in lines:
        if is_diagonal(line.start.x, line.start.y, line.end.x, line.end.y):
            diagonals.append(line)
    return diagonals


def make_grid(lines: list['Vent']) -> 'np.ndarray[int]':
    """ Get top left and bottom right points """
    x_values, y_values = [], []
    for line in lines:
        x_coord = [line.start.x, line.end.x]
        y_coord = [line.start.y, line.end.y]
        x_values.extend(x_coord)
        y_values.extend(y_coord)
    return np.zeros((max(y_values) + 1, max(x_values) + 1), dtype=int)


def mark_grid(grid: 'np.ndarray[int]', start: int, end: int, r: int,
              c: int) -> None:
    """ Update an array along rows or columns by one """
    for i in range(start, end):
        if r is not None and c is not None:
            # A hack to cancel out the negative value
            grid[abs(r)][c] += 1
            r += 1
            c += 1
        elif r is not None and c is None:
            grid[r][i] += 1
        elif c is not None and r is None:
            grid[i][c] += 1


def make_paths(grid: 'np.ndarray[int]', v: 'Vent') -> None:
    """ Set up paths from vents """
    # For diagonals
    if is_diagonal(v.start.x, v.start.y, v.end.x, v.end.y):
        initial_point = v.start.x
        end_point = v.end.x + 1
        row = v.start.y if v.start.y < v.end.y else -v.start.y
        col = v.start.x
    # For horizontals
    elif v.start.y == v.end.y:
        initial_point = v.start.x
        end_point = v.end.x + 1
        row, col = v.start.y, None
    # For verticals
    elif v.start.x == v.end.x:
        initial_point = v.start.y
        end_point = v.end.y + 1
        row, col = None, v.start.x
    mark_grid(grid, initial_point, end_point, row, col)


def part1(data: list[int]) -> int:
    """ Solve part 1 """
    eligible_vents = get_straights(data)
    grid = make_grid(eligible_vents)
    for vent in eligible_vents:
        make_paths(grid, vent)
    return np.count_nonzero(grid > 1)


def part2(data: list[int]) -> int:
    """ Solve part 2 """
    eligible_vents = get_straights(data) + get_diagonals(data)
    grid = make_grid(eligible_vents)
    for vent in eligible_vents:
        make_paths(grid, vent)
    return np.count_nonzero(grid > 1)


def solve(puzzle_input: str) -> tuple[int, int]:
    """ Solve the puzzle for the given input """
    data = parse(puzzle_input)
    solution1 = part1(data)  # Correct answer was 7085 (with my data)
    solution2 = part2(data)  # Correct answer was 20271 (with my data)

    return solution1, solution2


if __name__ == "__main__":
    for path in sys.argv[1:]:
        print(f"{path}:")
        puzzle_input = pathlib.Path(path).read_text().strip()
        solutions = solve(puzzle_input)
        print('\n'.join(str(solution) for solution in solutions))