aoc2024/day-09/internal/two/solve.go

package two

import (
	"bufio"
	"bytes"
	"io"
	"slices"
	"strconv"
)

const empty = -1

func Solve(buf []byte) (int, error) {
	r := bytes.NewReader(buf)
	expandedRaw, err := parseLines(r)
	if err != nil {
		return 0, err
	}

	disk := newDisk(expandedRaw)
	disk.defragment()
	disk.write()

	var sum int
	for i, n := range disk.data {
		if n == empty {
			continue
		}

		sum += i * n
	}

	return sum, nil
}

func parseLines(r io.Reader) ([]int, error) {
	var line string
	scanner := bufio.NewScanner(r)
	for scanner.Scan() {
		line = scanner.Text()
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	raw := [][]int{}
	for i, id := 0, 0; i < len(line); i, id = i+2, id+1 {
		raw = append(raw, []int{id, mustConv(string(line[i]))})
		var free int
		if i < len(line)-1 {
			free = mustConv(string(line[i+1]))
		}
		raw = append(raw, []int{empty, free})
	}

	expandedRaw := []int{}
	for _, vals := range raw {
		segment := slices.Repeat([]int{vals[0]}, vals[1])
		expandedRaw = append(expandedRaw, segment...)
	}

	return expandedRaw, nil
}

func mustConv(s string) int {
	n, err := strconv.Atoi(s)
	if err != nil {
		panic(err)
	}
	return n
}