2020/18

y2020/p10.py

+import collections
+import itertools
+import sys
+
+
+def recurse(n, sources, cache={0: 1}):
+    if n not in cache:
+        cache[n] = sum(recurse(m, sources) for m in sources[n])
+    return cache[n]
+
+
+text = sys.stdin.read()
+
+ns = [int(n) for n in text.splitlines()]
+ns += [0, max(ns) + 3]
+ns.sort()
+
+a, b = collections.Counter([b - a for a, b in zip(ns, ns[1:])]).values()
+print(a * b)
+
+sources = collections.defaultdict(list)
+for a, b in itertools.combinations(ns, 2):
+    if b - a <= 3:
+        sources[b].append(a)
+print(recurse(ns[-1], sources))