import copy
from collections import Counter


def load(filename):
	with open(filename, "r", encoding="utf-8") as f:
		r = sorted([int(x) for x in f])
		r.insert(0, 0)
		r.append(r[-1] + 3)
		return r


def calc_diffs(adapters):
	c = Counter()
	for i in range(len(adapters) - 1):
		d = adapters[i+1] - adapters[i]
		c[d] += 1
	return c


def calc_total(adapters):
	c = Counter()
	c[0] = 1
	for a in adapters:
		for i in range(1, 4):
			c[a] += c[a-i]
	return c[adapters[-1]]


adapters = load("test1.txt")
c = calc_diffs(adapters)
assert (c[1], c[3]) == (7, 5)
assert calc_total(adapters) == 8

adapters = load("test2.txt")
c = calc_diffs(adapters)
assert (c[1], c[3]) == (22, 10)
assert calc_total(adapters) == 19208

adapters = load("data.txt")
c = calc_diffs(adapters)
print(f"1: Answer is {c[1] * c[3]}")

c = calc_total(adapters)
print(f"2: Answer is {c}")