adventofcode2024/day05/sol.py

import functools

def parse(file):
    rules = []
    for line in file:
        line = line.strip()
        if line == "":
            break
        a, _, b = line.partition('|')
        rules.append((a,b))
    updates = []
    for line in file:
        line = line.strip()
        updates.append(line.split(','))
    return rules, updates

def solve(file):
    rules, updates = parse(file)
    rules = set(rules)
    def isvalid(pages):
        for i in range(len(pages)):
            p = pages[i]
            for g in pages[i+1:]:
                if (g,p) in rules:
                    return False
        return True

    # part 1
    t = 0
    invalid = []
    for pages in updates:
        if isvalid(pages):
            print(pages)
            t += int(pages[len(pages)//2])
        else:
            invalid.append(pages)
    print(t)

    # part 2: select applicable rules and form a graph
    # find the root node
    # do a topological sort by doing a DFS of the graph
    t = 0
    for pages in invalid:
        # select rules
        rules2 = [(p,g) for p,g in rules if p in pages and g in pages]
        #print(pages, rules2)

        # find root node
        S = set(p for p,g in rules2)
        E = set(g for p,g in rules2)
        top = next(iter(S-E))

        # topological sort
        l = []
        seen = set()
        def visit(node):
            if node in seen: return
            seen.add(node)
            for p,g in rules2:
                if p == node and g not in seen:
                    visit(g)
            l.append(node)
        visit(top)
        l.reverse()

        # sanity check
        assert len(l) == len(pages)
        assert isvalid(l)
        print(l)

        # add middle number to total
        t += int(l[len(l)//2])
    print(t)


solve(open("sample1.in"))
solve(open("input"))