advent-of-code/2024/python/day05.py

import matrix
import shared

def read(rows):
    second = False
    rules = []
    updates = []
    for row in rows:
        if second:
            updates.append(list(map(int, row.split(","))))
        elif row == "":
            second = True
        else:
            rules.append(list(map(int, row.split("|"))))
    return rules, updates

def get_bads(rules, updates):
    bads = []
    for idx, update in enumerate(updates):
        for first, second in rules:
            if first in update and second in update:    
                first_idx = update.index(first)
                second_idx = update.index(second)
                if first_idx > second_idx:
                    bads.append(idx)
                    break
    return bads

# @shared.profile
def part1(rows):
    rules, updates = read(rows)
    bads = get_bads(rules, updates)
    total = 0
    for idx, update in enumerate(updates):
        if idx in bads:
            continue
        total += update[len(update) // 2]
    print(total)

def fix_bad(rules, bad):
    for first, second in rules:
        if first in bad and second in bad: 
            first_idx = bad.index(first)
            second_idx = bad.index(second)
            if first_idx > second_idx:
                bad[first_idx], bad[second_idx] = bad[second_idx], bad[first_idx]
    return bad

# @shared.profile
def part2(rows):
    rules, updates = read(rows)
    bads = get_bads(rules, updates)
    total = 0
    for idx in bads:
        bad = updates[idx]
        while True:
            after = fix_bad(rules, bad[:])
            if bad == after:
                break
            bad = after
        total += bad[len(bad) // 2]
    print(total)


def main():
    rows = [row for row in shared.load_rows(5)]
    with shared.elapsed_timer() as elapsed:
        part1(rows)
        print("🕒", elapsed())

    rows = [row for row in shared.load_rows(5, True)]
    with shared.elapsed_timer() as elapsed:
        part2(rows)
        print("🕒", elapsed())


if __name__ == "__main__":
    main()