当青训营遇上码上掘金，会擦出怎样的火花？

「青训营 X 码上掘金」主题创作活动入营版

在今年的第五届字节跳动青训营中，新引入了“青豆”的积分制度。而在「青训营 X 码上掘金」主题创作活动中，也出现了青豆相关的主题，这引起了我强烈的兴趣。所以我选择了主题 4：攒青豆。攒青豆是一道关于栈、数组、双指针、动态规划、单调栈的题目，较为困难，十分具有挑战性。

题目

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

输入：height = [5,0,2,1,4,0,1,0,3]

输出：17

解析：上面是由数组 [5,0,2,1,4,0,1,0,3] 表示的柱子高度，在这种情况下，可以接 17 个单位的青豆。

解答

我使用 Java 语言完成了这个主题，先给出代码，之后进行解释说明。

我的题解 - 码上掘金

import java.util.*;

public class Main {
    public static void main(String[] args) {
        /* 读入数据，以逗号分割 */
        Scanner sc = new Scanner(System.in);
        String str = sc.nextLine();
        String[] arr = str.split(",");
        int[] height = new int[arr.length];
        for (int i = 0; i < arr.length; i++) {
            height[i] = Integer.parseInt(arr[i]);
        }

        /* 输出青豆数量 */
        System.out.println(trap(height));
    }

    /* 计算青豆数量 */
    public static int trap(int[] height) {
        int n = height.length;
        int[] preMax = new int[n], sufMax = new int[n];
        preMax[0] = height[0];
        for (int i = 1; i < n; ++i) {
            preMax[i] = Math.max(preMax[i - 1], height[i]);
        }
        sufMax[n - 1] = height[n - 1];
        for (int i = n - 2; i >= 0; --i) {
            sufMax[i] = Math.max(sufMax[i + 1], height[i]);
        }
        int ans = 0;
        for (int i = 0; i < n; ++i) {
            ans += Math.min(preMax[i], sufMax[i]) - height[i];
        }
        return ans;
    }
}

在我本机的测试结果如下图所示，可以看到正确输出了答案。

解释说明

接下来对我的代码进行一些说明。

可以把每个宽度为1的区域看作为一个广义的桶，桶能容纳的青豆数量首先取决于它所在区域的柱子高度，也取决于左右两侧所有区域的柱子高度的最大值。

比如说在这个例子中，下标为2的区域中柱子长度为2，它左边的区域为[5,0]，右边的区域为[1,4,0,1,0,3]。显然，该区域的左边区域中柱子高度最大值为5，右边区域中柱子高度最大值为4。也就是说，这是一个底高为2的“桶”，“桶”的“左木板”的高度为5，“右木板”的高度为4。所以，这个桶中能容纳的青豆数量为：min(5,4) - 2 = 2 个单位。

所有区域，即所有桶中能够容纳的青豆数量加起来的总和，即为题目所求的量。

可以遍历每个桶，再分别计算它左右“两个木板”的高度，得出青豆数量。但这样的代码显然有些臃肿，时间复杂度太高，不是优秀的解法。

我的做法是，先求出木板高度，再依次计算青豆数量，减小重复计算。我使用了preMax和sufMax两个数组，储存每个区域对应的左右两块木板的最大高度。

步骤如下：

读入数据，以逗号分割

Scanner sc = new Scanner(System.in);
String str = sc.nextLine();
String[] arr = str.split(",");
int[] height = new int[arr.length];
for (int i = 0; i < arr.length; i++) {
    height[i] = Integer.parseInt(arr[i]);
}

定义前缀数组、后缀数组

int n = height.length;
int[] preMax = new int[n], sufMax = new int[n];

计算前缀最大高度，即左木板最大高度

preMax[0] = height[0];
for (int i = 1; i < n; ++i) {
    preMax[i] = Math.max(preMax[i - 1], height[i]);
}

计算后缀最大高度，即右木板最大高度

sufMax[n - 1] = height[n - 1];
for (int i = n - 2; i >= 0; --i) {
    sufMax[i] = Math.max(sufMax[i + 1], height[i]);
}

计算各个区域能容纳的青豆总和

int ans = 0;
for (int i = 0; i < n; ++i) {
    ans += Math.min(preMax[i], sufMax[i]) - height[i];
}
return ans;

输出青豆数量

System.out.println(trap(height));

攒青豆！

题目

解答