aoc2024/day9/main.go

package main

import (
	"bufio"
	"fmt"
	"os"
	"strings"
)

const SENTINEL int = -2
const FREE_BLOCK int = -1

type Block struct {
	Prev, Next *Block
	Size       int
	FileId     int
}

func (before *Block) InsertAfter(after *Block) {
	before.Next.Prev = after
	after.Next, before.Next = before.Next, after
	after.Prev = before
}

func (block *Block) Detach() {
	block.Prev.Next, block.Next.Prev = block.Next, block.Prev
	block.Prev, block.Next = nil, nil
}

func (block *Block) TryMergeFreeWithNext() {
	if block.FileId == FREE_BLOCK && block.Next.FileId == FREE_BLOCK {
		block.Size += block.Next.Size
		block.Next.Detach()
	}
}

func (block *Block) TryMergeFree() *Block {
	block.TryMergeFreeWithNext()
	if block.Prev.FileId == FREE_BLOCK {
		block = block.Prev
		block.TryMergeFreeWithNext()
	}
	return block
}

func (block *Block) Split(Size int) {
	if Size >= block.Size {
		panic("please")
	}

	if Size == block.Size {
		return
	}

	newBlock := &Block{nil, nil, block.Size - Size, block.FileId}
	block.Size = Size
	block.InsertAfter(newBlock)
}

type Disk struct {
	Sentinel *Block
}

func readInt(b byte) int {
	return int(b - '0')
}

func MakeDisk(description string) (disk Disk) {
	disk.Sentinel = new(Block)
	disk.Sentinel.Size = 0
	disk.Sentinel.FileId = SENTINEL
	disk.Sentinel.Next = disk.Sentinel
	disk.Sentinel.Prev = disk.Sentinel

	nextFileId := 0
	nextIsFree := false

	head := disk.Sentinel

	for _, b := range []byte(description) {
		block := &Block{
			Prev: nil, Next: nil, Size: readInt(b),
		}

		if block.Size == 0 {
			nextIsFree = !nextIsFree
			continue
		}

		if nextIsFree {
			block.FileId = FREE_BLOCK
			nextIsFree = false
		} else {
			block.FileId = nextFileId
			nextFileId++
			nextIsFree = true
		}

		head.InsertAfter(block)
		head = block
	}

	return
}

func (disk *Disk) Print() {
	fmt.Println("===============")
	for b := disk.Sentinel.Next; b != disk.Sentinel; b = b.Next {
		fmt.Println(b)
	}
	fmt.Println("===============")
	bufio.NewReader(os.Stdin).ReadBytes('\n')
}

func (disk *Disk) Compactify() {
	head := disk.Sentinel.Next
	tail := disk.Sentinel.Prev

	for head != tail {
		if head.FileId >= 0 {
			head = head.Next
			continue
		}

		if tail.FileId == FREE_BLOCK {
			tail = tail.Prev
			continue
		}

		if tail.Size > head.Size {
			tail.Split(tail.Size - head.Size)
			tail = tail.Next
		}

		if head.Size > tail.Size {
			head.Split(tail.Size)
		}

		head.FileId, tail.FileId = tail.FileId, head.FileId
		tail = tail.TryMergeFree()
	}
}

func (disk *Disk) Defragment() {
	tail := disk.Sentinel.Prev

	for tail != disk.Sentinel {
		if tail.FileId == FREE_BLOCK {
			tail = tail.Prev
			continue
		}

		var freeBlock *Block = nil
		for head := disk.Sentinel.Next; head != tail; head = head.Next {
			if head.FileId == FREE_BLOCK && head.Size >= tail.Size {
				freeBlock = head
				break
			}
		}

		if freeBlock == nil {
			tail = tail.Prev
			continue
		}

		if freeBlock.Size > tail.Size {
			freeBlock.Split(tail.Size)
		}

		freeBlock.FileId, tail.FileId = tail.FileId, freeBlock.FileId
		tail = tail.TryMergeFree()
	}
}

func (disk *Disk) ComputeChecksum() (sum int64) {
	pos := 0
	for b := disk.Sentinel.Next; b != disk.Sentinel; b = b.Next {
		if b.FileId >= 0 {
			for i := 0; i < b.Size; i++ {
				sum += int64(pos * b.FileId)
				pos++
			}
		} else {
			pos += b.Size
		}
	}
	return
}

func main() {
	disk := MakeDisk(readData("example.txt"))
	disk.Compactify()
	fmt.Println(disk.ComputeChecksum(), 1928)

	disk = MakeDisk(readData("data.txt"))
	disk.Compactify()
	fmt.Println(disk.ComputeChecksum(), 6415184586041)

	disk = MakeDisk(readData("example.txt"))
	disk.Defragment()
	fmt.Println(disk.ComputeChecksum(), 2858)

	disk = MakeDisk(readData("data.txt"))
	disk.Defragment()
	fmt.Println(disk.ComputeChecksum())
}

func readData(fileName string) string {
	resultBytes, err := os.ReadFile(fileName)
	if err != nil {
		panic(err)
	}

	return strings.TrimSpace(string(resultBytes))
}