about summary refs log tree commit diff
path: root/aoc/2024/02/part-two.py
blob: ad5f5a67484b1617dbd59c1dd8063246ffea3740 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
from collections import Counter
from itertools import chain, count
from sys import stdin


def safe(levels, dampened=False):
    head = levels[0 if 0 < levels[1] - levels[0] < 4 else 1] - 1,
    tail = levels[-1 if 0 < levels[-1] - levels[-2] < 4 else -2] + 1,
    dampenable = tuple((not dampened and 0 < a - c < 4
                        and safe(levels[:i]+levels[i+1:], True))
                       for i, a, b, c in zip(count(),
                                             chain(levels[1:], tail),
                                             levels,
                                             chain(head, levels))
                       if not (0 < a - b < 4 and 0 < b - c < 4))
    return not dampenable or any(dampenable)


print(sum(safe(levels) or safe(tuple(map(int.__neg__, levels)))
          for levels in (tuple(map(int, line))
                         for line in map(str.split, stdin))))