def read_file(): with open("input.txt","r") as f: return [int(x) for x in f.read().split('\n')] def part1(adapters): jolts = [0] + sorted(adapters) # j(port) = 0 d1 = 0 d3 = 1 # j(device) = max + 3 i = 0 while i <= len(adapters): difference = jolts[i] - jolts[i-1] if difference == 1: d1 += 1 elif difference == 2: d2 += 1 elif difference == 3: d3 += 1 i += 1 return d1 * d3 """ def find_all_paths(graph, start, end, path=[]): path = path + [start] if start == end: return [path] if start not in graph: return [] paths = [] for node in graph[start]: if node not in path: newpaths = find_all_paths(graph, node, end, path) for newpath in newpaths: paths.append(newpath) return paths """ def part2(adapters): jolts = [0] + sorted(adapters) from collections import Counter dp = Counter() dp[0] = 1 jolts = sorted(adapters) jolts.append(jolts[-1] + 3) for j in jolts: dp[j] = dp[j-1] + dp[j-2] + dp[j-3] return dp[jolts[-1]] """ graph = {} for i, joltage in enumerate(jolts): graph[joltage] = [x for x in jolts[i+1:i+4] if 0 < x - joltage <= 3] # scan next 3, assuming diff of +1 +1 +1 at most return len( find_all_paths( graph, 0, max(jolts) ) ) """ def main(): adapters = read_file() print(part1(adapters), part2(adapters)) if __name__ == "__main__": main()