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