「青训营 X 码上掘金」攒青豆

79 阅读3分钟

「青训营 X 码上掘金」攒青豆

当青训营遇上码上掘金,青训营公布了一道名为“攒青豆”的题目,是一道非常经典的算法题,下面来讲讲我的解题思路。

题目描述

现有 n 个宽度为 1 的柱子,给出 n 个非负整数依次表示柱子的高度,排列后如下图所示,此时均匀从上空向下撒青豆,计算按此排列的柱子能接住多少青豆。(不考虑边角堆积)

方法一:动态规划

对于下标 ii,撒青豆后能到达的最大高度等于下标 ii 两边的最大高度的最小值,下标 ii 处能接的青豆数量等于下标 ii 处的青豆能到达的最大高度减去 height[i]\textit{height}[i]

朴素的做法是对于数组 height\textit{height} 中的每个元素,分别向左和向右扫描并记录左边和右边的最大高度,然后计算每个下标位置能接的青豆量。假设数组 height\textit{height} 的长度为 nn,该做法需要对每个下标位置使用 O(n)O(n) 的时间向两边扫描并得到最大高度,因此总时间复杂度是 O(n2)O(n^2)

上述做法的时间复杂度较高是因为需要对每个下标位置都向两边扫描。如果已经知道每个位置两边的最大高度,则可以在 O(n)O(n) 的时间内得到能接的青豆总量。使用动态规划的方法,可以在 O(n)O(n) 的时间内预处理得到每个位置两边的最大高度。

创建两个长度为 nn 的数组 leftMax\textit{leftMax}rightMax\textit{rightMax}。对于 0i<n0 \le i<nleftMax[i]\textit{leftMax}[i] 表示下标 ii 及其左边的位置中,height\textit{height} 的最大高度,rightMax[i]\textit{rightMax}[i] 表示下标 ii 及其右边的位置中,height\textit{height} 的最大高度。

显然,leftMax[0]=height[0]\textit{leftMax}[0]=\textit{height}[0]rightMax[n1]=height[n1]\textit{rightMax}[n-1]=\textit{height}[n-1]。两个数组的其余元素的计算如下:

  • 1in11 \le i \le n-1 时,leftMax[i]=max(leftMax[i1],height[i])\textit{leftMax}[i]=\max(\textit{leftMax}[i-1], \textit{height}[i])
  • 0in20 \le i \le n-2 时,rightMax[i]=max(rightMax[i+1],height[i])\textit{rightMax}[i]=\max(\textit{rightMax}[i+1], \textit{height}[i])

因此可以正向遍历数组 height\textit{height} 得到数组 leftMax\textit{leftMax} 的每个元素值,反向遍历数组 height\textit{height} 得到数组 rightMax\textit{rightMax} 的每个元素值。

在得到数组 leftMax\textit{leftMax}rightMax\textit{rightMax} 的每个元素值之后,对于 0i<n0 \le i<n,下标 ii 处能接的青豆量等于 min(leftMax[i],rightMax[i])height[i]\min(\textit{leftMax}[i],\textit{rightMax}[i])-\textit{height}[i]。遍历每个下标位置即可得到能接的青豆总量。

class Solution:
    def trap(self, height: List[int]) -> int:
        if not height:
            return 0
        
        n = len(height)
        leftMax = [height[0]] + [0] * (n - 1)
        for i in range(1, n):
            leftMax[i] = max(leftMax[i - 1], height[i])

        rightMax = [0] * (n - 1) + [height[n - 1]]
        for i in range(n - 2, -1, -1):
            rightMax[i] = max(rightMax[i + 1], height[i])

        ans = sum(min(leftMax[i], rightMax[i]) - height[i] for i in range(n))
        return ans

方法二:单调栈

除了计算并存储每个位置两边的最大高度以外,也可以用单调栈计算能接的青豆总量。

维护一个单调栈,单调栈存储的是下标,满足从栈底到栈顶的下标对应的数组 height\textit{height} 中的元素递减。

从左到右遍历数组,遍历到下标 ii 时,如果栈内至少有两个元素,记栈顶元素为 top\textit{top}top\textit{top} 的下面一个元素是 left\textit{left},则一定有 height[left]height[top]\textit{height}[\textit{left}] \ge \textit{height}[\textit{top}]。如果 height[i]>height[top]\textit{height}[i]>\textit{height}[\textit{top}],则得到一个可以接青豆的区域,该区域的宽度是 ileft1i-\textit{left}-1,高度是 min(height[left],height[i])height[top]\min(\textit{height}[\textit{left}],\textit{height}[i])-\textit{height}[\textit{top}],根据宽度和高度即可计算得到该区域能接的青豆量。

为了得到 left\textit{left},需要将 top\textit{top} 出栈。在对 top\textit{top} 计算能接的青豆量之后,left\textit{left} 变成新的 top,重复上述操作,直到栈变为空,或者栈顶下标对应的\textit{top},重复上述操作,直到栈变为空,或者栈顶下标对应的 \textit{height}中的元素大于或等于中的元素大于或等于\textit{height}[i]$

在对下标 ii 计算能接的青豆量之后,将 ii 入栈,继续遍历后面的下标,计算能接的青豆量。遍历结束之后即可得到能接的青豆总量。

class Solution:
    def trap(self, height: List[int]) -> int:
        ans = 0
        stack = list()
        n = len(height)
        
        for i, h in enumerate(height):
            while stack and h > height[stack[-1]]:
                top = stack.pop()
                if not stack:
                    break
                left = stack[-1]
                currWidth = i - left - 1
                currHeight = min(height[left], height[i]) - height[top]
                ans += currWidth * currHeight
            stack.append(i)
        
        return ans