Advent of Code 2020 Day 10

"""Day 10 of Advent of Code 2020 Solution"""
from collections import defaultdict


def adapter_io(file_location) -> list[int]:
    with open(file_location, "r") as f:
        data = [int(line.strip()) for line in f.readlines()]
        data = sorted(data)
        data = [0] + data + [data[-1] + 3]
    return data


def n_connection_dist(data: list[int]) -> tuple[int, int]:
    adapter_differences = defaultdict(int)
    for i, rating in enumerate(data[1:], 1):
        adapter_differences[rating - data[i - 1]] += 1
    return (adapter_differences[1], adapter_differences[3])


def part_a(file_location) -> int:
    data = adapter_io(file_location)
    _1, _3 = n_connection_dist(data)
    return _1 * _3


def part_b(file_location):
    data = adapter_io(file_location)
    graph_nodes = defaultdict(list)
    for i, n in enumerate(data):
        next_actual_nodes = set(data[i+1:i+4])
        next_possible_nodes = set([n + i for i in range(1, 4)])
        for intersecting_node in next_possible_nodes.intersection(next_actual_nodes):
            graph_nodes[n].append(intersecting_node)

    n_paths_table = {max(data): 1}
    for node in reversed(data[:-1]):
        n_paths = 0
        for next_node in graph_nodes[node]:
            n_paths += n_paths_table[next_node]
        n_paths_table[node] = n_paths

    return n_paths_table[min(data)]


if __name__ == '__main__':
    file_location = r"data\day10.txt"
    print(part_a(file_location))
    print(part_b(file_location))