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()))