aoc-2020/15.py

lines = []

with open("15.input") as f:
    for line in f.readlines():
        lines.append(line.strip())

numbers = [int(num) for num in lines[0].split(",")]

def find_nth(numbers, end):
    numbers = list(numbers)
    last_spoken = {}

    for (i, num) in enumerate(numbers[:-1]):
        last_spoken[num] = i

    last = numbers[-1]
    numbers = numbers[0:-1]

    for i in range(len(numbers), end):
        if last in last_spoken:
            next = i - last_spoken[last]
        else:
            next = 0
        numbers.append(last)
        last_spoken[last] = i
        last = next

    return numbers[-1]

print("Answer 1:", find_nth(numbers, 2020))
print("Answer 2:", find_nth(numbers, 30000000))