import aocd

def parse_reports(data):
    for line in data.split('\n'):
        yield [int(level) for level in line.split()]

def is_safe(orig_report):
    modified = [orig_report[:i] + orig_report[i+1:] for i in range(len(orig_report))]
    for report in [orig_report] + modified:
        pairs = list(zip(report[:-1], report[1:]))
        diffs = [a-b for a,b in pairs]
        if all(diff in (1, 2, 3) for diff in diffs) or all(diff in (-1, -2, -3) for diff in diffs):
            return True
    return False

def solve(data):
    return sum(is_safe(report) for report in parse_reports(data))

if __name__ == '__main__':
    solution = solve(aocd.get_data(year=2024, day=2))

    print(solution)
    aocd.submit(solution, year=2024, day=2, part='b')