day05/solution_gemini-1.5-pro.py

from collections import defaultdict

def parse_input(file):
    with open(file, 'r') as f:
        lines = f.read().splitlines()
    rules_raw, updates_raw = lines[:lines.index('')], lines[lines.index('') + 1:]
    rules = {}
    for rule in rules_raw:
        a, b = map(int, rule.split('|'))
        rules[(a, b)] = True
    updates = []
    for update in updates_raw:
        updates.append(list(map(int, update.split(','))))
    return rules, updates

def is_correct_order(update, rules):
    for i in range(len(update)):
        for j in range(i + 1, len(update)):
            if (update[j], update[i]) in rules:
                return False
    return True

def get_middle_page(update):
    return update[len(update) // 2]

def topological_sort(update, rules):
    graph = defaultdict(list)
    in_degree = defaultdict(int)
    nodes = set(update)
    for i in nodes:
        for j in nodes:
            if (i, j) in rules:
                graph[i].append(j)
                in_degree[j] += 1

    queue = [node for node in nodes if in_degree[node] == 0]
    result = []
    while queue:
        u = queue.pop(0)
        result.append(u)
        for v in graph[u]:
            in_degree[v] -= 1
            if in_degree[v] == 0:
                queue.append(v)
    return result

file = "./input.txt"
rules, updates = parse_input(file)

correct_sum = 0
incorrect_sum = 0

for update in updates:
    if is_correct_order(update, rules):
        correct_sum += get_middle_page(update)
    else:
        correct_update = topological_sort(update, rules)
        incorrect_sum += get_middle_page(correct_update)

print(correct_sum)
print(incorrect_sum)