aoc2024/day09.py

#!/usr/bin/env python3
from copy import copy
from typing import Optional
from dataclasses import dataclass


@dataclass
class Block:
    fid: Optional[int]
    length: int


puzzle = []
with open("day09") as f:
    puzzle = list(map(int, f.readline().strip()))

fs = []

def fill_fs():
    fid = 0
    for i,e in enumerate(puzzle):
        if i%2==0: # file
            fs.append(Block(fid, e))
            fid += 1
        else: # gap
            fs.append(Block(None, e))

def find_gaps():
    gaps = []
    for i,b in enumerate(fs):
        if b.fid is None: gaps.append((i, b))
    return gaps

def find_files():
    files = []
    for i,b in enumerate(fs):
        if b.fid is not None: files.append((i, b))
    return files

def checksum():
    checksum = 0
    idx = 0
    for b in fs:
        if b.fid is None:
            for i in range(b.length): idx += 1
        else:
            for i in range(b.length):
                checksum += b.fid*idx
                idx += 1

    return checksum

def part1():
    while True:
        gaps = find_gaps()
        if not gaps: break

        if gaps[-1][0] == len(fs)-1: # free block at end
            del fs[-1]
            continue

        files = find_files()
        next_gap = gaps[0][0]
        next_file = files[-1][0]

        if fs[next_gap].length == fs[next_file].length:
            fs[next_gap] = fs[next_file]
            del fs[-1]
            continue

        if fs[next_gap].length > fs[next_file].length:
            gap_length = fs[next_gap].length
            file_length = fs[next_file].length

            fs.insert(next_gap, fs[next_file])
            fs[next_gap+1].length = gap_length-file_length
            del fs[-1]

            continue

        if fs[next_gap].length < fs[next_file].length:
            gap_length = fs[next_gap].length
            file_length = fs[next_file].length
            fid = fs[next_file].fid

            fs[next_file].length = file_length - gap_length
            fs[next_gap] = Block(fid=fid, length=gap_length)

            continue

fill_fs()
part1()
print(checksum())

def part2():
    files = list(reversed(find_files()))

    for i,file in enumerate(files):
        gaps = find_gaps()
        if not gaps: break

        file_idx = fs.index(file[1])

        for gap in gaps:
            if gap[0] >= file_idx:
                break

            if gap[1].length == file[1].length:
                fs[file_idx] = Block(fid=None, length=file[1].length)
                fs[gap[0]] = file[1]
                break

            if gap[1].length > file[1].length:
                gap_length = gap[1].length
                file_length = file[1].length

                fs[file_idx] = Block(fid=None, length=file[1].length)

                fs.insert(gap[0], file[1])
                fs[gap[0]+1].length = gap_length-file_length
                break

fill_fs()
part2()
print(checksum())