aoc24-day9b

# Setup

import sys
import time

sys.path.append('../utils')
from pyutils import *

sample="2333133121414131402"

with open('input.txt', 'r', encoding='utf-8') as f:
    puzzle = f.read()

def read_diskmap(diskmap: str) -> list[int]:
    processed: list[int] = []
    mode: int = 1 # 1 == file, -1 == free space
    files: int = -1
    for digit in diskmap:
        if not digit.isnumeric():
            continue
        digit = int(digit)
        if mode == 1:
            files += 1
            processed.extend([files] * digit)
        elif mode == -1:
            processed.extend([-1] * digit)
        mode *= -1
    return processed

def block_report(disk: list[int]):
    groups = {'used': {}, 'free': []}
    """
    gruops dict structure is: {
        'used': {
            <int; file id>: (<int; index of first block for this file>, [<list[int]; block contents>])
        },
        'free': [
            (<int; index of first block for this free space>, <int; size of this free space, i.e. number of blocks>)
        ]
    }
    """
    freeblock = 0
    for n, file_id in enumerate(disk):
        if file_id == -1:
            freeblock += 1
        else:
            if freeblock > 0:
                # If we've reached a non-free block, and we *were* counting free blocks before, save the counted blocks
                groups['free'].append((n - freeblock, freeblock))
                freeblock = 0
            if file_id not in groups['used']:
                groups['used'][file_id] = (n, [])
            groups['used'][file_id][1].append(file_id)
    return groups

def optimize_disk(disk: list[int], defrag: bool=True) -> list[int]:
    disk = disk.copy()
    if defrag:
        report = block_report(disk)
        # Reverse order of used blocks to iterate from the end as per instructions
        report['used'] = {k:report['used'][k] for k in sorted(report['used'].keys(), reverse=True)}
        for used_id, ublock in report['used'].items():
            # Tried to visualize what was going on so I could fix it, but it didn't help much...
            # clear_output(wait=True)
            # time.sleep(0.1)
            # print(''.join(map(str, d)).replace('-1', '.')[:1000])
            for free_id, fblock in enumerate(report['free']):
                if fblock[1] >= len(ublock[1]):
                    fblock = report['free'].pop(free_id)
                    # Replace free space with file
                    disk[fblock[0]:fblock[0] + len(ublock[1])] = ublock[1]
                    # Replace previous file space with free space
                    disk[ublock[0]:ublock[0] + len(ublock[1])] = [-1] * len(ublock[1])
                    # Add free block back into same list position, adjust index and shrink size in report
                    report['free'].insert(free_id, (fblock[0] + len(ublock[1]), fblock[1] - len(ublock[1])))
                    break
    else:
        for n in range(len(disk)):
            n = abs(n - len(disk)) - 1
            if disk[n] != -1:
                leftmost_free = disk.index(-1)
                if leftmost_free > n:
                    # All sorted, no more work to do
                    break
                disk[leftmost_free], disk[n] = disk[n], disk[leftmost_free]
    return disk

def calc_disk_checksum(d: list[int], show: bool=False) -> int:
    chsum: int = 0
    for n, digit in enumerate(d):
        if show:
            time.sleep(0.005)
            clear_output(wait=True)
            print(f'{repr(n):^8}{repr(digit):^8}{chsum}')
        if digit == -1:
            continue
        chsum += n * digit
    return chsum

# Solve

disk: list[int] = read_diskmap(puzzle)

ta = time.perf_counter()
optimized = optimize_disk(disk, defrag=True)
tb = time.perf_counter()
print(f'Optimized disk of length {len(optimized)} in {tb - ta:.05f}s')

calc_disk_checksum(optimized, show=False)