AoC 2021 Day 15 in Go

package main

import (
    "aoc/utils"
    "fmt"
)

func main() {
    input := utils.ReadLines()
    risk, width, height := parseGrid(input)
    fmt.Println("part 1: ", FindBestPath(risk, width, height))
    risk, width, height = parseExtendedGrid(input)
    fmt.Println("part 2: ", FindBestPath(risk, width, height))
}

type state struct {
    risk   []uint16 // table of risks
    best   []uint16 // best[i] = cost of the lowest *known* path to i
    width  int
    height int
    length int
}

// FindBestPath returns the cost of the best path from the top-left to
// the bottom-right.
//
// This algorithm assumes that the best path will move right or down most of the time.
// We first estimate best(x, y) to be:
//    min(best(x-1, y), best(x, y-1)) + risk(x, y)
//
// Once this initial guess is done, for each already estimated neighbor (xn, yn),
// we check that:
//    best(xn, yn) <= best(x, y) + risk(xn, yn)
//
// If not, it means that we have found a shorter path to (xn, yn) that involves
// going up or left, and so we fix best(xn, yn) and propagate the check
// recursively.
func FindBestPath(risk []uint16, width, height int) uint16 {
    s := state{
        risk:   risk,
        best:   make([]uint16, len(risk)),
        width:  width,
        height: height,
        length: len(risk),
    }
    s.initFirstRow()
    s.computeBestPaths()
    return s.best[s.length-1]
}

func (s *state) initFirstRow() {
    for i := 1; i < s.width; i++ {
        s.best[i] = s.best[i-1] + s.risk[i]
    }
}

func (s *state) computeBestPaths() {
    for y := 1; y < s.height; y++ {
        i := y * s.width
        s.best[i] = s.best[i-s.width] + s.risk[i]
        for x := 1; x < s.width; x++ {
            i++
            s.best[i] = s.lowestNeighbor(i) + s.risk[i]
            s.fixPaths(i)
        }
    }
}

func (s *state) lowestNeighbor(i int) uint16 {
    return min(s.best[i-1], s.best[i-s.width])
}

func (s *state) fixPaths(i int) {
    currentBest := s.best[i]
    s.forEachNeighbor(i, func(j int) {
        if b := currentBest + s.risk[j]; b < s.best[j] {
            // we found a better path to this neighbor, let's propagate it
            s.best[j] = b
            s.fixPaths(j)
        }
    })
}

func (s *state) forEachNeighbor(i int, loop func(int)) {
    x := i % s.width
    if i > s.width {
        loop(i - s.width)
    }
    if x != 0 {
        loop(i - 1)
    }
    if x < s.width-1 {
        loop(i + 1)
    }
    if j := i + s.width; j < s.length {
        loop(i + s.width)
    }
}

func parseGrid(input []string) ([]uint16, int, int) {
    height, width := len(input), len(input[0])
    grid := make([]uint16, width*height)
    for y, line := range input {
        for x, num := range line {
            grid[y*width+x] = uint16(num - '0')
        }
    }
    return grid, width, height
}

func parseExtendedGrid(input []string) ([]uint16, int, int) {
    height, width := 5*len(input), 5*len(input[0])
    grid := make([]uint16, width*height)
    for y, line := range input {
        for x, num := range line {
            for i := uint16(0); i < 5; i++ {
                yy := y + int(i)*len(input)
                for j := uint16(0); j < 5; j++ {
                    xx := x + int(j)*len(input[0])
                    grid[yy*width+xx] = wrap(uint16(num-'0') + i + j)
                }
            }
        }
    }
    return grid, width, height
}

func wrap(i uint16) uint16 {
    for i > 9 {
        i -= 9
    }
    return i
}

func min(a, b uint16) uint16 {
    if a < b {
        return a
    }
    return b
}