AoC 2022 Day 8

# aoc202208.py

import pathlib
import sys

import numpy as np

# Typing aliases
M = list[list[int]]


def parse(puzzle_input: str) -> M:
    """ Parse input """
    return [[int(x) for x in row] for row in puzzle_input.splitlines()]


def mark_visible(arr, tree) -> None:
    """ Mark matching trees in given array """
    max_height = arr[0]
    for i, v in enumerate(arr[:-1]):
        if max_height < arr[i + 1]:
            max_height = arr[i + 1]
            if tree[i + 1] == 0:
                tree[i + 1] = 1


def add_trees(arr, grid) -> None:
    """ Mark matching trees in grid """
    rows, cols = arr.shape
    grid[0] = [1 for n in range(cols)]
    grid[-1] = [1 for n in range(cols)]
    for n in range(1, rows - 1):
        mark_visible(arr[n], grid[n])
        mark_visible(arr[n][::-1], grid[n][::-1])


def visible_trees(heights):
    """ Return grid of matching trees """
    matches = np.zeros(heights.shape, dtype=int)
    add_trees(heights, matches)
    add_trees(np.transpose(heights), np.transpose(matches))
    return matches


def part1(data: M) -> int:
    """ Solve part 1 """
    tree_heights = np.array(data)
    return np.count_nonzero(visible_trees(tree_heights))


def find_candidates(matrix):
    """ Return coordinates of peaks """
    coords: list[tuple[int, int]] = []
    last_index: int = matrix.shape[1] - 1
    for coord in np.argwhere(matrix):
        # Remove edges
        if not any([coord[0] in (0, last_index), coord[1] in (0, last_index)]):
            coords.append([coord[0], coord[1]])
    return coords


def view_count(value: int, span: list[int], level: bool = False) -> int:
    """ Return score of visible trees """
    count: int = 0
    for n in span:
        if level:
            continue
        count += 1
        if n == value:
            level = True
        if n > value:
            break
    return count


def scenic_score(arr: list, idx: int) -> tuple[int, int]:
    """ Return scenic score for one tree """
    v: int = arr[idx]
    to_start: list[int] = arr[:idx][::-1]
    to_end: list[int] = arr[idx + 1:]
    return view_count(v, to_start), view_count(v, to_end)


def get_scores(matrix, coords) -> list[int]:
    """ get a lits of scenic scores for all candidate trees """
    scores: list[int] = []
    flipped = np.transpose(matrix)
    for row, col in coords:
        # left, right, up, down are ints
        left, right = scenic_score(matrix[row], col)
        up, down = scenic_score(flipped[col], row)
        scores.append(np.prod([up, left, down, right]))
    return scores


def part2(data: M) -> int:
    """ Solve part 2 """
    tree_heights = np.array(data)
    candidate_trees: list[tuple[int, int]] = find_candidates( visible_trees(tree_heights))
    scores: list[int] = get_scores(tree_heights, candidate_trees)
    return max(scores)


def solve(puzzle_input: str) -> tuple[int, int]:
    """ Solve the puzzle for the given input """
    data: M = parse(puzzle_input)
    solution1: int = part1(data)  # Correct answer was 1717 (with my data)
    solution2: int = part2(data)  # Correct answer was 321975 (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))