/ Python / 2020 / 10.py
10.py
 1  from lib import *
 2  
 3  input = read_input(2020, 10)
 4  
 5  from collections import Counter
 6  
 7  nums = sorted(ints(input))
 8  
 9  
10  prev = 0
11  counter = Counter()
12  for x in nums:
13      d = x - prev
14      counter[d] += 1
15      prev = x
16  
17  counter[3] += 1
18  print(counter[1] * counter[3])
19  
20  
21  dp = {}
22  
23  
24  def count(idx, joltage):
25      if idx == len(nums):
26          return joltage == nums[-1]
27  
28      if nums[idx] - joltage > 3:
29          return 0
30  
31      if (idx, joltage) not in dp:
32          dp[(idx, joltage)] = count(idx + 1, joltage) + count(idx + 1, nums[idx])
33  
34      return dp[(idx, joltage)]
35  
36  
37  print(count(0, 0))