package q3651

import (
	"math"
	"slices"
)

func minCost(grid [][]int, k int) int {
	w, h := len(grid[0]), len(grid)
	gridVal := func(coord int) int { return grid[coord/w][coord%w] }

	sortedCoords := make([]int, w*h)
	for i := range sortedCoords {
		sortedCoords[i] = (i)
	}
	slices.SortFunc(sortedCoords, func(a, b int) int { return gridVal(a) - gridVal(b) })

	cache := make([]int32, w*h)
	prevMinCosts := make([]int32, len(sortedCoords))

	for tele := range k + 1 {
		// Precompute min costs
		if tele > 0 {
			var curMin int32 = math.MaxInt32
			i, j := 0, 0
			for i < len(sortedCoords) {
				curMin = min(curMin, cache[sortedCoords[i]])

				i++
				for j < i && (i == len(sortedCoords) || gridVal(sortedCoords[i]) > gridVal(sortedCoords[j])) {
					prevMinCosts[sortedCoords[j]] = curMin
					j++
				}
			}
		}

		for r := h - 1; r >= 0; r-- {
			for c := w - 1; c >= 0; c-- {
				if r == h-1 && c == w-1 {
					continue
				}

				var val int32 = math.MaxInt32
				if r+1 < h {
					val = min(val, cache[(r+1)*w+c]+int32(grid[r+1][c]))
				}
				if c+1 < w {
					val = min(val, cache[r*w+c+1]+int32(grid[r][c+1]))
				}

				if tele == 0 {
					cache[r*w+c] = val
					continue
				}

				// Find min cost with teleports
				cache[r*w+c] = min(val, prevMinCosts[r*w+c])
			}
		}
	}

	return int(cache[0])
}

var _ = minCost