https://leetcode.cn/problems/trapping-rain-water/?envType=study-plan-v2&envId=top-interview-150

go

package leetcode150

import (
    "sort"
    "testing"
)

func TestTrap(t *testing.T) {
    height := []int{4, 2, 0, 3, 2, 5}
    res := trap3(height)
    println(res)
}
// 官方题解三。双指针
func trap5(height []int) (ans int) {
    left, right := 0, len(height)-1
    leftMax, rightMax := 0, 0
    for left < right {
       leftMax = max(leftMax, height[left])
       rightMax = max(rightMax, height[right])
       if height[left] < height[right] {
          ans += leftMax - height[left]
          left++
       } else {
          ans += rightMax - height[right]
          right--
       }
    }
    return
}

func max(a, b int) int {
    if a > b {
       return a
    }
    return b
}

// 官方题解一，动态规划
func trap3(height []int) int {
    if len(height) < 2 {
       return 0
    }
    leftMax := make([]int, len(height))
    rightMax := make([]int, len(height))
    leftMax[0] = height[0]
    rightMax[len(height)-1] = height[len(height)-1]
    for i := 1; i < len(height); i++ {
       if leftMax[i-1] < height[i] {
          leftMax[i] = height[i]
       } else {
          leftMax[i] = leftMax[i-1]
       }
    }
    for i := len(height) - 2; i >= 0; i-- {
       if rightMax[i+1] < height[i] {
          rightMax[i] = height[i]
       } else {
          rightMax[i] = rightMax[i+1]
       }
    }
    total := 0
    for i, h := range height {
       if i-1 >= 0 && i+1 < len(height) {
          wall := 0
          if leftMax[i] < rightMax[i] {
             wall = leftMax[i]
          } else {
             wall = rightMax[i]
          }
          if wall > h {
             total += wall - h
          }
       }
    }
    return total
}

// 反向思维，计算总容量，减去不能接水的部分,时间复杂度O(nlog(n))，空间复杂度O(n)
func trap2(height []int) int {
    if len(height) < 2 {
       return 0
    }

    wallTotal := 0
    heightPosMap := map[int][]int{}
    for i, h := range height {
       // 只需要计算最两侧的墙壁
       if len(heightPosMap[h]) <= 1 {
          heightPosMap[h] = append(heightPosMap[h], i)
       } else {
          heightPosMap[h][1] = i
       }

       wallTotal += h
    }

    var heightArr []int
    for h := range heightPosMap {
       heightArr = append(heightArr, h)
    }
    sort.Sort(sort.Reverse(sort.IntSlice(heightArr)))

    width := len(height)
    // 计算总容量，并减去墙壁
    allTotal := heightArr[0]*width - wallTotal
    for i, h := range heightArr {
       // 做减法
       if len(heightPosMap[h]) < 2 {
          // 只有一个墙壁的情况
          if i+1 < len(heightArr) {
             hd := heightArr[i] - heightArr[i+1]
             allTotal = allTotal - hd*width + hd
          }
       } else {
          // 两侧有墙壁
          if i+1 < len(heightArr) {
             hd := heightArr[i] - heightArr[i+1]
             allTotal -= heightPosMap[h][0]*hd + (width-heightPosMap[h][1]-1)*hd
          }
       }
       // 维护最两侧的墙壁
       if i+1 < len(heightArr) {
          nextH := heightArr[i+1]
          if len(heightPosMap[h]) < 2 && len(heightPosMap[nextH]) < 2 {
             if heightPosMap[h][0] > heightPosMap[nextH][0] {
                heightPosMap[nextH] = append(heightPosMap[nextH], heightPosMap[h][0])
             } else {
                heightPosMap[nextH] = append(heightPosMap[nextH], heightPosMap[h][0])
                heightPosMap[nextH][1], heightPosMap[nextH][0] = heightPosMap[nextH][0], heightPosMap[nextH][1]
             }
          } else if len(heightPosMap[h]) < 2 && len(heightPosMap[nextH]) == 2 {
             if heightPosMap[h][0] < heightPosMap[nextH][0] {
                heightPosMap[nextH][0] = heightPosMap[h][0]
             } else if heightPosMap[h][0] > heightPosMap[nextH][1] {
                heightPosMap[nextH][1] = heightPosMap[h][0]
             }
          } else if len(heightPosMap[nextH]) < 2 && len(heightPosMap[h]) == 2 {
             if heightPosMap[h][1] > heightPosMap[nextH][0] {
                heightPosMap[nextH] = append(heightPosMap[nextH], heightPosMap[h][1])
                if heightPosMap[h][0] < heightPosMap[nextH][0] {
                   heightPosMap[nextH][0] = heightPosMap[h][0]
                }
             } else {
                heightPosMap[nextH] = append(heightPosMap[nextH], heightPosMap[h][0])
                heightPosMap[nextH][1], heightPosMap[nextH][0] = heightPosMap[nextH][0], heightPosMap[nextH][1]
             }
          } else {
             if heightPosMap[nextH][0] > heightPosMap[h][0] {
                heightPosMap[nextH][0] = heightPosMap[h][0]
             }
             if heightPosMap[nextH][1] < heightPosMap[h][1] {
                heightPosMap[nextH][1] = heightPosMap[h][1]
             }
          }
       }
    }
    return allTotal
}

// 超时未过（测试用例过了322/323），正向思维，计算能接多少水
func trap(height []int) int {
    heightPosMap := map[int][]int{}
    for i, h := range height {
       heightPosMap[h] = append(heightPosMap[h], i)
    }

    var heightArr []int
    for h := range heightPosMap {
       heightArr = append(heightArr, h)
    }

    sort.Sort(sort.Reverse(sort.IntSlice(heightArr)))

    total := 0
    for i, h := range heightArr {
       // 说明能有两侧墙壁
       if len(heightPosMap[h]) >= 2 {
          if i+1 < len(heightArr) {
             hd := heightArr[i] - heightArr[i+1]
             posArr := heightPosMap[h]
             sort.Sort(sort.IntSlice(posArr))
             for j := range posArr {
                if j+1 < len(posArr) {
                   // -1 是减去墙壁
                   total += (posArr[j+1] - posArr[j] - 1) * hd
                }
             }
          }
       }

       if i+1 < len(heightArr) {
          nextH := heightArr[i+1]
          heightPosMap[nextH] = append(heightPosMap[nextH], heightPosMap[h]...)
       }
    }
    return total
}