记一道美团面试编程题的解题过程：从暴力到优化
By Long Luo
最近在做 LeetCode学习计划 中的 2021届秋季校招笔试真题 时，在这道题上 meituan-002. 小美的仓库整理 卡了一点时间，在此记录下自己从困惑到最终解决的全过程，吸取经验教训，提高自己知识水平。
题目
这是一道 Medium 难度的题目，如下所示：
meituan-002. 小美的仓库整理
小美是美团仓库的管理员，她会根据单据的要求按顺序取出仓库中的货物，每取出一件货物后会把剩余货物重新堆放，使得自己方便查找。已知货物入库的时候是按顺序堆放在一起的。如果小美取出其中一件货物，则会把货物所在的一堆物品以取出的货物为界分成两堆，这样可以保证货物局部的顺序不变。已知货物最初是按 1~n 的顺序堆放的，每件货物的重量为 w[i] ，小美会根据单据依次不放回的取出货物。请问根据上述操作，小美每取出一件货物之后，重量和最大的一堆货物重量是多少？
格式：
输入：
- 输入第一行包含一个正整数 n ，表示货物的数量。
- 输入第二行包含 n 个正整数，表示 1~n 号货物的重量 w[i] 。
- 输入第三行有 n 个数，表示小美按顺序取出的货物的编号，也就是一个1~n的全排列。
输出：
- 输出包含n行，每行一个整数，表示每取出一件货物以后，对于重量和最大的一堆货物，其重量和为多少。示例：
输入：
5
3 2 4 4 5
4 3 5 2 1
输出：
9
5
5
3
0解释： 原本的状态是[[3, 2, 4, 4, 5]]，取出4号货物后，得到 [[3,2,4], [5]]，第一堆货物的和是9，然后取出3号货物得到[[3, 2] [5]]，此时第一堆和第二堆的和都是5，以此类推。
提示： 1 <= n <= 50000 1 <= w[i] <= 100
请注意，本题需要自行编写「标准输入」和「标准输出」逻辑，以及自行import/include需要的library。了解书写规则
从题意来看，这道题目并不难，所以当时看到题目之后，马上刷刷开始做题，觉得太简单了！
解题思路
首先想到的肯定是暴力法。
暴力法
根据题意，就是求数组某个区间的区间和。求区间和最简单的方式当然是遍历了，单次遍历复杂度为 \(O(N)\) ，而全部结果计算的时间复杂度为 \(O(N^2)\) ，于是就有了下面的这个解法：1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38public static void main(String[] args) {
Scanner scanner = new Scanner(System.in);
int n = scanner.nextInt();
scanner.nextLine();
int[] weight = new int[n];
for (int i = 0; i < n; i++) {
weight[i] = scanner.nextInt();
}
int[] number = new int[n];
for (int i = 0; i < n; i++) {
number[i] = scanner.nextInt();
}
int[] ans = new int[n];
for (int i = 0; i < n; i++) {
weight[number[i] - 1] = 0;
ans[i] = getMax(weight, number[i]);
}
for (int i = 0; i < n; i++) {
System.out.println(ans[i]);
}
}
private static int getMax(int[] weight, int no) {
int leftTotal = 0;
int rightTotal = 0;
for (int i = 0; i < (no - 1); i++) {
leftTotal += weight[i];
}
for (int i = no; i < weight.length; i++) {
rightTotal += weight[i];
}
return Math.max(leftTotal, rightTotal);
}
这个答案只通过了示例中的测试用例，其他的都没有通过，那么问题出在什么地方呢？
通过重新读题，再加上题目给的示例的迷惑下，我明白了错在什么地方了！
我以为只需要每次拿掉一个货物之后，以此货物为界，分成左右 2 堆，然后返回更大的一堆的值，而且示例中连续拿了相邻的货物，也让我产生了错觉，因为错的离谱！
这可是一道中等题啊！肯定不是这么简单的啊！
重新完成了正确的暴力解决法如下：1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47public static void main(String[] args) {
Scanner scanner = new Scanner(System.in);
int n = scanner.nextInt();
scanner.nextLine();
int[] weight = new int[n];
for (int i = 0; i < n; i++) {
weight[i] = scanner.nextInt();
}
int[] number = new int[n];
for (int i = 0; i < n; i++) {
number[i] = scanner.nextInt();
}
int[] ans = new int[n];
for (int i = 0; i < n; i++) {
weight[number[i] - 1] = 0;
ans[i] = getMax(weight);
}
for (int i = 0; i < n; i++) {
System.out.println(ans[i]);
}
}
private static int getMax(int[] weight) {
int maxSum = 0;
int n = weight.length;
int idx = 0;
while (idx < n) {
while (idx < n && weight[idx] == 0) {
idx++;
}
if (idx < n && weight[idx] != 0) {
int sum = 0;
while (idx < n && weight[idx] != 0) {
sum += weight[idx];
idx++;
}
maxSum = Math.max(sum, maxSum);
}
}
return maxSum;
}
因为上述时间复杂度为 \(O(N^2)\) ，铁定超时，所以需要寻找更优的解决方案！
前缀和
很明显暴力法具有很大的优化空间，因为很多求和操作都是重复的。因此我们将货物进行预处理一下，使用前缀和数组便可做到花费 \(O(1)\) 的时间来查询区间和了！
TreeSet
快速查询区间和做到了，已知前缀和，求区间和还需要什么呢？
对了，左边界与右边界。
以示例中的第1次取货来说：[3, 2, 4, 4, 5] 。取了4号货物后，数组将以4号位置为轴心，分割为左区间 [3, 2, 4] 和右区间 [5] 。
前缀和数组分割后对应为：[0, 3, 5, 9] [13] [18]。
此时，左区间中的左边界为 0，右边界为 3 ，可得 leftSum = 9 - 0 = 9。右区间中的左边界为 4 ，右边界为 5 ，可得 rightSum = 18 - 13 = 5。
此时，可以快速插入与查找的有序容器 Set 派上了用场。
Set 将完成以下工作：查询分割点的左右边界。
在完成分割工作后，将分割点加入其中，因为其将成为后续分割点查询所在的子数组的边界。
TreeMap
在完成分割与分割点左右区间和的计算后，需要将左右区间和放入容器中，然后查询容器中区间和的最大值。
若使用无序容器，则每次查找最大值需要消耗 \(O(N)\) 的时间，这样的话前面的优化算是白做了。
因此，我们将所有区间和放入有序容器中进行查询，此时需要考虑 多个区间和的值 相等的问题，于是我们自然而然地想到了Map。
Map将完成以下工作：
- 首先将包含分割点的区间和
SegSum的数量减1，因为它将被分割点拆开，若数量为 1 ，则直接删除； - 将分割后得到的 左右区间和加入其中；
- 返回容器中的最大 key ，其为当前场上存在的 最大区间和。
代码如下所示：1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54public static void main(String[] args) {
Scanner scanner = new Scanner(System.in);
int n = scanner.nextInt();
scanner.nextLine();
int[] weight = new int[n];
for (int i = 0; i < n; i++) {
weight[i] = scanner.nextInt();
}
int[] number = new int[n];
for (int i = 0; i < n; i++) {
number[i] = scanner.nextInt();
}
int[] prefixSum = new int[n + 1];
for (int i = 0; i < n; i++) {
prefixSum[i + 1] = weight[i] + prefixSum[i];
}
int[] ans = new int[n];
TreeSet<Integer> boundSet = new TreeSet<>();
boundSet.add(0);
boundSet.add(n + 1);
TreeMap<Integer, Integer> sumMap = new TreeMap<>();
for (int i = 0; i < n; i++) {
int pos = number[i];
int left = boundSet.lower(pos);
int right = boundSet.higher(pos);
boundSet.add(pos);
int segSum = prefixSum[right - 1] - prefixSum[left];
Integer count = sumMap.get(segSum);
if (count != null) {
if (count == 1) {
sumMap.remove(segSum);
} else {
sumMap.put(segSum, count - 1);
}
}
int leftSum = prefixSum[pos - 1] - prefixSum[left];
int rightSum = prefixSum[right - 1] - prefixSum[pos];
sumMap.put(leftSum, sumMap.getOrDefault(leftSum, 0) + 1);
sumMap.put(rightSum, sumMap.getOrDefault(rightSum, 0) + 1);
ans[i] = sumMap.lastKey();
}
for (int i = 0; i < n; i++) {
System.out.println(ans[i]);
}
}
- 时间复杂度：\(O(N \logN)\) ，其中 \(N\) 来自 \(n\) 次询问，有序容器的查找、插入和删除操作均摊为 \(\logN\) 。
- 空间复杂度：\(O(N)\) 。
逆向思考
我们不妨逆向操作，一开始仓库里没有货物，我们按小美抽取货物的顺序倒序插入。
那么，初始化最大重量 maxWeight 为 0 ，我们每次插入后，只会改变插入点下标的左右相邻两堆货物，将它们连成一堆货物。
因此，只需要求出新产生的货物堆的重量之和，更新 maxWeight 。
如果采取遍历的方式获取当前插入点的连通区域，时间复杂度是 \(O(N^2)\) 会超时，对于如何快速获取我们可以使用一个 \(2n\) 的数组 bounds，用于记录一个堆货物的左右边界下标。
对于这道题，每次操作必然不会在重复的位置。因此，只需要维护一堆货物的左边界和右边界，无需关心中间区域，每次操作的位置一定相邻货物或远离货物，不会存在重叠。而对于求任意区间内的和，我们可以使用前缀和数组实现 \(O(1)\) 的时间复杂度。
如果
bounds[x + 1][0] != -1，说明右边有一堆相邻的货物，将右边界right更新为bounds[x + 1][1]，并使当前新产生的堆重量 cur 加上 prev[bounds[x + 1][1]] - prev[bounds[x + 1][0] - 1]如果
bounds[x - 1][1] != -1，说明左边有一堆相邻的货物，将左边界left更新为bounds[x - 1][0]，并使当前新产生的堆重量 cur 加上 prev[bounds[x - 1][1]] - prev[bounds[x - 1][0] - 1]
例如：
5 3 2 4 4 5 4 3 5 2 1
这个例子
初始化 maxW = 0 , bounds [[-1, -1], [-1, -1], [-1, -1], [-1, -1], [-1, -1], [-1, -1]],
前缀和数组 prev={0, 3, 5, 9, 13, 18},
原始货物 weight={0, 3, 5, 9, 13, 18}
第一次操作： x为1，先记录一次maxW(因为是倒着操作的，先记录再做改变)
- 初始化新产生的一堆货物的左边界 left = x，右边界 right = x ，重量 cur = arr[1] = 3 。
- 判断左右有无一堆相邻的货物
- 得到cur = 3，更新maxW = max(maxW, cur) = 3,
- 更新新产生的堆的左边界bounds[left][0] = left;bounds[left][1] = right;
- 更新堆的左边界 bounds[right][0] = left;bounds[right][1] = right;
第二次操作： x为 2 ，先记录一次 maxW
- 初始化新产生的一堆货物的左边界 left = 2 ，右边界 right = 2，重量 cur = 2 。
- 然后判断，由于 bounds[x + 1][0] != -1 ,说明左边有一堆相邻的货物，cur += 3，left = 1。
- 得到cur = 5，更新 maxW = max(maxW,cur) = 5,
- 更新新产生的堆的左边界bounds[left][0] = left; bounds[left][1] = right;
- 更新堆的左边界 bounds[right][0] = left; bounds[right][1] = right;
之后操作同理可得。
代码如下所示：1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54public static void main(String[] args) {
Scanner scanner = new Scanner(System.in);
int n = Integer.parseInt(scanner.nextLine());
String[] weights = scanner.nextLine().split(" ");
int[] weight = new int[n + 1];
int[] prefixSum = new int[n + 1];
for (int i = 1; i <= n; i++) {
weight[i] = Integer.parseInt(weights[i - 1]);
prefixSum[i] = prefixSum[i - 1] + weight[i];
}
String[] number = scanner.nextLine().split(" ");
int[][] bounds = new int[n + 1][2];
for (int i = 0; i <= n; i++) {
bounds[i] = new int[]{-1, -1};
}
int[] ans = new int[n];
int maxWeight = 0;
for (int i = n - 1; i >= 0; i--) {
int x = Integer.parseInt(number[i]);
ans[i] = maxWeight;
if (i == 0) {
break;
}
// 更新最大重量
int cur = weight[x];
int left = x;
int right = x;//左边界和右边界,注意如果左右无连通区域则区间为[x,x],所以初始化为x
//每次只会将左右两块区域连成一块,我们只需关心一段区间的左边界和右边界,就能通过前缀和数组查询到区间和
if (x + 1 <= n && bounds[x + 1][0] != -1) {
cur += prefixSum[bounds[x + 1][1]] - prefixSum[bounds[x + 1][0] - 1];
right = bounds[x + 1][1]; //更新右边界
}
if (x - 1 > 0 && bounds[x - 1][1] != -1) {
cur += prefixSum[bounds[x - 1][1]] - prefixSum[bounds[x - 1][0] - 1];
left = bounds[x - 1][0]; //更新左边界
}
maxWeight = Math.max(maxWeight, cur);
// 更新两端点的左右区间
bounds[left][0] = left;
bounds[left][1] = right;
bounds[right][0] = left;
bounds[right][1] = right;
}
for (int i = 0; i < n; i++) {
System.out.println(ans[i]);
}
}
总结
通过这道题，学习到了几个道理：
- 首先必须读懂题意，不要搞错题意，这是前提；
- 先使用最直观的暴力法，看能否解决问题；
- 暴力法不可行的话，想一想如何用空间换时间，减少时间复杂度；
- 学会逆向思考，想一想还有没有其他方法可以解决这个问题，提高自己融会贯通的能力。