aoc2020day10.py

def part2(data):
    def candidates(n):
        return 1 + (n * (n + 1) // 2)
    lines = [int(v) for v in data.split("\n")]
    nums = list(sorted(lines)) + [max(lines) + 3]
    a = 0
    i = 0
    total = 1
    while i < len(nums):
        b = nums[i]
        diff = b - a
        if diff != 3:
            found = False
            for j, lookahead in enumerate(zip(nums[i:], nums[i+1:]), start=1):
                if not lookahead[1] - lookahead[0] != diff:
                    continue
                if j > 1:
                    total *= candidates(lookahead[0] - a - 1)
                    a = lookahead[0]
                    i += j
                    found = True
                break
            if found:
                continue
        j = 1
        a = b
        i += j
    return total