advent_of_code_2024/day2/run.py

from copy import deepcopy
from typing import List

EXAMPLE = """
7 6 4 2 1
1 2 7 8 9
9 7 6 2 1
1 3 2 4 5
8 6 4 4 1
1 3 6 7 9
"""


class CheckLevels:
    """Checks that levels are safe"""

    def __init__(self, input: str):
        self.levels: List[List[int]] = [
            [int(val) for val in line.split()] for line in input.splitlines() if line
        ]

    def __str__(self):
        return f"Part 1 = {self._part1}, Part 2 = {self._part2}"

    @property
    def _part1(self) -> int:
        return len([val for val in self.get_safe(problem_damper=False) if val])

    @property
    def _part2(self) -> int:
        return len([val for val in self.get_safe(problem_damper=True) if val])

    @staticmethod
    def get_differences(line) -> List[List[int]]:
        return [val - line[i] for i, val in enumerate(line[1:])]

    def get_safe(self, problem_damper: bool = False) -> List[List[bool]]:
        def is_level_safe(line: List[int]) -> bool:
            all_same = all([val < 0 for val in line]) or all([val > 0 for val in line])
            all_in_range = (
                min([abs(val) for val in line]) >= 1
                and max([abs(val) for val in line]) <= 3
            )
            return all_same and all_in_range

        outp = []
        for line in self.levels:
            safe = is_level_safe(self.get_differences(line))
            if problem_damper and not safe:
                for i in range(len(line)):
                    reduced = deepcopy(line)
                    reduced.pop(i)
                    reduced_safe = is_level_safe(self.get_differences(reduced))
                    if reduced_safe:
                        safe = reduced_safe
                        break
            outp.append(safe)
        return outp


if __name__ == "__main__":
    example = CheckLevels(input=EXAMPLE)
    print("EXAMPLE:\n", example)
    assert example._part1 == 2
    assert example._part2 == 4
    with open("input.txt", "r") as f:
        print("Puzzle:\n", CheckLevels(input=f.read()))