leetcode-go/hard/407.trapping-rain-water-ii.go

func trapRainWater(heightMap [][]int) (water int) {
	m := len(heightMap)
	if m == 0 {
		return
	}
	n := len(heightMap[0])
	if n == 0 {
		return
	}
	visited := make([][]bool, m)
	for i := 0; i < m; i++ {
		visited[i] = make([]bool, n)
	}
	var pq PriorityQueue
	for i := 0; i < m; i++ {
		heap.Push(&pq, []int{heightMap[i][0], 0, i})
		heap.Push(&pq, []int{heightMap[i][n-1], n - 1, i})
	}
	for i := 1; i < n-1; i++ {
		heap.Push(&pq, []int{heightMap[0][i], i, 0})
		heap.Push(&pq, []int{heightMap[m-1][i], i, m - 1})
	}
	max := 0
	for pq.Len() != 0 {
		cell := heap.Pop(&pq).([]int)
		if max < cell[0] {
			max = cell[0]
		}
		for i := 0; i < 4; i++ {
			nx, ny := cell[1]+dx[i], cell[2]+dy[i]
			if 0 < nx && nx < n-1 && 0 < ny && ny < m-1 && !visited[ny][nx] {
				h := heightMap[ny][nx]
				if h < max {
					water += max - h
				}
				heap.Push(&pq, []int{h, nx, ny})
				visited[ny][nx] = true
			}
		}
	}
	return
}

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

type PriorityQueue [][]int

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

func (pq *PriorityQueue) Push(x interface{}) {
	*pq = append(*pq, x.([]int))
}

func (pq *PriorityQueue) Pop() interface{} {
	x := (*pq)[pq.Len()-1]
	*pq = (*pq)[:pq.Len()-1]
	return x
}