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

发表于 2021-09-10 更新于 2022-05-01 分类于 Data Structures and Algorithms 阅读次数：本文字数： 7.7k 阅读时长 ≈ 7 分钟

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^2)$ ，于是就有了下面的这个解法：

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堆，然后返回更大的一堆的值，而且示例中连续拿了相邻的货物，也让我产生了错觉，因为错的离谱！

这可是一道中等题啊！肯定不是这么简单的啊！

重新完成了正确的暴力解决法如下：

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(N^2)$ ，铁定超时，所以需要寻找更优的解决方案！

2.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，其为当前场上存在的最大区间和。

代码如下所示：

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)$ ，其中 $N$ 来自 $n$ 次询问，有序容器的查找、插入和删除操作均摊为 $logN$ 。
空间复杂度： $O(N)$ 。

2.3 逆向思考

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

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

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

如果采取遍历的方式获取当前插入点的连通区域，时间复杂度是 $O(N^2)$ 会超时，对于如何快速获取我们可以使用一个 $n*2$ 的数组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;

之后操作同理可得。

代码如下所示：

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

总结

通过这道题，学习到了几个道理：

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