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)