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 }