记一道美团面试编程题的解题过程:从暴力到优化

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。了解书写规则

从题意来看,这道题目并不难,所以当时看到题目之后,马上刷刷开始做题,觉得太简单了!

解决思路

首先想到的肯定是暴力法

2.1 暴力法

根据题意,就是求数组某个区间的区间和。求区间和最简单的方式当然是遍历了,复杂度为O(N)O(N),而全部结果计算的时间复杂度为O(N2)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
38
public 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
47
public 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(N2)O(N^2),铁定超时,所以需要寻找更优的解决方案!

2.2 前缀和

很明显暴力法具有很大的优化空间,因为很多求和操作都是重复的。因此我们将货物进行预处理一下,使用前缀和数组便可做到花费O(1)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)O(N)的时间,这样的话前面的优化算是白做了。

因此,我们将所有区间和放入有序容器中进行查询,此时需要考虑 多个区间和的值 相等的问题,于是我们自然而然地想到了Map

Map将完成以下工作:

  1. 首先将包含分割点的区间和SegSum的数量减1,因为它将被分割点拆开,若数量为 1,则直接删除;
  2. 将分割后得到的 左右区间和加入其中;
  3. 返回容器中的最大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
54
public 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(NlogN)O(NlogN),其中NN来自nn次询问,有序容器的查找、插入和删除操作均摊为logNlogN

空间复杂度:O(N)O(N)

2.3 逆向思考

我们不妨逆向操作,一开始仓库里没有货物,我们按小美抽取货物的顺序倒序插入。

那么,初始化最大重量maxWeight为0,我们每次插入后,只会改变插入点下标的左右相邻两堆货物,将它们连成一堆货物。

因此,只需要求出新产生的货物堆的重量之和,更新maxWeight

如果采取遍历的方式获取当前插入点的连通区域,时间复杂度是O(N2)O(N^2)会超时,对于如何快速获取我们可以使用一个n2n*2的数组bounds,用于记录一个堆货物的左右边界下标。

对于这道题,每次操作必然不会在重复的位置。因此,只需要维护一堆货物的左边界和右边界,无需关心中间区域,每次操作的位置一定相邻货物或远离货物,不会存在重叠。而对于求任意区间内的和,我们可以使用前缀和数组实现O(1)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;

之后操作同理可得。

代码如下所示:

    public 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]);
        }
    }

总结

通过这道题,学习到了几个道理:

  1. 首先必须读懂题意,不要搞错题意,这是前提;
  2. 先使用最直观的暴力法,看能否解决问题;
  3. 暴力法不可行的话,想一想如何用空间换时间,减少时间复杂度;
  4. 学会逆向思考,想一想还有没有其他方法可以解决这个问题,提高自己融会贯通的能力。