LeetCode 42. Trapping Rain Water

傳送門

題意

給你$n$個非負整數，表示圍欄的高度（寬度都爲1），問你如果下雨一共能儲多少水？如圖：

Input: [0,1,0,2,1,0,1,3,2,1,2,1]
Output: 6

思路

Dp

可以發現對於每一個$i$來說，它能儲水的高度顯然取決於它左面和右面最高的圍欄高度之間的最小值。所以我們可以簡單的想到brute force，對每個$i$,分別找它左面和右面的最大值，然後答案就是$min(max_l[i],max_r[i]) - height[i]$。顯然這種方法的時間複雜度爲$O(n^2)$,所以我們可以通過$Dp$的思想將時間複雜度優化到$O(n)$級別。

class Solution:
    def trap(self, height):
        if not height:
            return 0
        _len = len(height)
        max_l = [height[i] for i in range(_len)]
        max_r = [height[i] for i in range(_len)]
        for i in range(1,_len):
            max_l[i] = max(max_l[i - 1],height[i])
        ans = 0
        for i in range(_len - 2,-1,-1):
            max_r[i] = max(max_r[i + 1],height[i])
        for i in range(_len):
            ans += min(max_l[i],max_r[i]) - height[i]
        return ans 

Two Pointer

上面的Dp的方法雖然是$O(n)$,但是仍然遍歷了很多次，我們可以使用Two pointer來進一步優化。
我們可以發現對於每個$i$來說，如果有$maxleft[i] < maxright[i]$,那麼$i$處的儲水量和$maxright[i]$就沒什麼關係了。因此，我們可以得出:

• 如果有$maxleft[i] < maxright[i]$,我們在這個過程中維護$maxleft$即可，如果$height[i] > maxleft$ 則更新，答案爲$ans += maxleft - height[i]$,$l += 1$ 繼續從左往右
• 如果有$maxleft[i] < maxright[i]$，那麼應該從右往左了，並維護$maxright$,答案爲$ans += maxright-height[i],r -= 1$

class Solution:
    def trap(self, height):
        if not height:
            return 0
        _len = len(height)
        max_l,max_r = 0,0
        l,r = 0,_len - 1
        ans = 0
        while l <= r:
            if max_l < max_r:
                max_l = max(max_l,height[l])
                ans += max_l - height[l]
                l += 1
            else:
                max_r = max(max_r,height[r])
                ans += max_r - height[r]
                r -= 1
        return ans

Stack

以上兩個方法都是對每個$i$找到左右最深能儲多少水，這裏其實我們還可以換個思想。對於每一個$i$,只要左右有比他高的，那麼他肯定能儲一部分水，如果有再高的又能多儲一部分——所以這裏我們的思想是可以一層層的解決儲水的問題。
那麼這裏我們就需要維護一個遞減棧，棧中保存下標。
如果有$height[i] > height[stack[top]]$,那麼說明stack[top]處可以儲水，離它左右最近的一層就是當前的$i$和棧中的下一個(棧頂)。那麼有:

• 這一層的長度爲$distance = i - stack[top]-1$
• 最大高度爲$h = min(height[i],height[stack[-1]]$
• 則儲水爲$ans += h \times distance$

class Solution:
   def trap(self, height):
       if not height:
           return 0
       ans = 0
       _len = len(height)
       stack = []
       for i in range(_len):
           while len(stack) and height[i] > height[stack[-1]]:
               top = stack.pop()
               if not len(stack):
                   break
               dis = i - stack[-1] - 1
               ans += dis * (min(height[i],height[stack[-1]]) - height[top])
           stack.append(i)
       return ans