leetcode-go/leetcode/hard/675.cut-off-trees-for-golf-event.go

type minHeap [][3]int

func (h minHeap) Len() int           { return len(h) }
func (h minHeap) Less(i, j int) bool { return h[i][0] < h[j][0] }
func (h minHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *minHeap) Push(x interface{}) {
	*h = append(*h, x.([3]int))
}

func (h *minHeap) Pop() interface{} {
	l := h.Len()
	x := (*h)[l-1]
	*h = (*h)[:l-1]
	return x
}

func cutOffTree(forest [][]int) (step int) {
	var h minHeap
	H, W := len(forest), len(forest[0])
	for y := range forest {
		for x, v := range forest[y] {
			if 1 < v {
				heap.Push(&h, [3]int{v, y, x})
			}
		}
	}
	beg := []int{0, 0}
	used := make([][]bool, H)
	for h.Len() != 0 {
		end := heap.Pop(&h).([3]int)
		for i := range used {
			used[i] = make([]bool, W)
		}
		ds := walk(forest, H, W, beg, end[1:], used)
		if ds == -1 {
			return -1
		}
		step += ds
		beg = end[1:]
	}
	return step
}

var dy = []int{0, 1, 0, -1}
var dx = []int{1, 0, -1, 0}

func walk(forest [][]int, H, W int, beg, end []int, used [][]bool) (step int) {
	queue := [][]int{beg}
	n := 1
	for len(queue) != 0 {
		pos := queue[0]
		queue = queue[1:]
		if pos[0] == end[0] && pos[1] == end[1] {
			return
		}
		for i := 0; i < 4; i++ {
			ny, nx := pos[0]+dy[i], pos[1]+dx[i]
			if 0 <= ny && ny < H && 0 <= nx && nx < W && !used[ny][nx] && forest[ny][nx] != 0 {
				queue = append(queue, []int{ny, nx})
				used[ny][nx] = true
			}
		}
		n--
		if n == 0 {
			step++
			n = len(queue)
		}
	}
	return -1
}