update

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

@@ -1,3 +1,61 @@
-func trapRainWater(heightMap [][]int) int {
-	
+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
 }