OA算法题：求最大子数组大小之和，现有代码仅过样例求正确解法

题目要求 #

给定长度为n的正整数数组arr：
对于数组每个下标i（0 ≤ i ≤ n-1），找到包含该元素的最大子数组（即最长子数组），使该元素为子数组的最大值，将子数组大小累加到初始为0的结果中。

示例 #

• 示例1：arr = [3,5,6]，结果为6（1+2+3=6）
• 示例2：arr = [1,2,1]，结果为5（1+3+1=5）
• 示例3：arr = [1,1,1,1]，结果为16（4×4=16）

我的错误代码（仅通过样例，隐藏用例全失败） #

public class Main {

  public static int solve(int[] arr) {
    int n = arr.length;
    int result = 0;
    int currentMax = Integer.MIN_VALUE;

    int size = 0;
    for (int i = 0; i < n; i++) {
      currentMax = Math.max(currentMax, arr[i]);
          
      if (arr[i] == currentMax) {
          size++;
          int backUp = size;
          for(int j=i+1; j<n; j++) {
              if(arr[j] <= currentMax) {
                  size++;
              } else {
                  break;
              }
          }
        result += size;
        size = backUp;
      } else {
          size = 1;
          currentMax = arr[i];
          result++;
      }
    }

    return result;
  }

  public static void main(String[] args) {
    System.out.println(solve(new int[]{3,5,6}));//6
    System.out.println(solve(new int[]{1,2,1}));//5
    System.out.println(solve(new int[]{1,1,1,1}));//16
  }
}

问题分析 #

你的代码核心错误在于：

1. 错误维护全局currentMax和size变量，将前序遍历的状态带入当前元素的计算，无法针对每个元素独立求解；
2. 仅考虑了当前元素向右延伸的子数组范围，完全忽略向左延伸的可能，导致无法找到真正的最长符合条件的子数组。

正确解法：单调栈 #

要高效解决问题，我们需要为每个元素找到：

• 左边第一个比它大的元素的下标（记为left[i]，无则为-1）；
• 右边第一个比它大的元素的下标（记为right[i]，无则为数组长度n）。

包含当前元素的最长符合条件的子数组，就是从left[i]+1到right[i]-1的区间，长度为right[i] - left[i] - 1，累加所有元素的该长度即可得到结果。

正确Java代码 #

import java.util.Stack;

public class Main {

    public static int solve(int[] arr) {
        int n = arr.length;
        int[] left = new int[n]; // 左边第一个大于arr[i]的下标，无则为-1
        int[] right = new int[n]; // 右边第一个大于arr[i]的下标，无则为n
        Stack<Integer> stack = new Stack<>();

        // 计算left数组
        for (int i = 0; i < n; i++) {
            while (!stack.isEmpty() && arr[stack.peek()] <= arr[i]) {
                stack.pop();
            }
            left[i] = stack.isEmpty() ? -1 : stack.peek();
            stack.push(i);
        }

        stack.clear();

        // 计算right数组
        for (int i = n - 1; i >= 0; i--) {
            while (!stack.isEmpty() && arr[stack.peek()] <= arr[i]) {
                stack.pop();
            }
            right[i] = stack.isEmpty() ? n : stack.peek();
            stack.push(i);
        }

        // 累加所有元素对应的最长子数组长度
        int result = 0;
        for (int i = 0; i < n; i++) {
            result += right[i] - left[i] - 1;
        }
        return result;
    }

    public static void main(String[] args) {
        System.out.println(solve(new int[]{3,5,6})); // 6
        System.out.println(solve(new int[]{1,2,1})); //5
        System.out.println(solve(new int[]{1,1,1,1})); //16
        System.out.println(solve(new int[]{1,3,2,3})); // 10（正确结果）
    }
}

代码说明 #

1. 单调栈求left数组：遍历数组时，栈中保存下标对应的元素单调递减。弹出栈中所有值<=当前元素的下标后，栈顶即为左边第一个更大元素的下标，否则为-1，最后将当前下标入栈。
2. 单调栈求right数组：从右往左遍历，同样维护单调递减栈，找到右边第一个更大元素的下标，否则为数组长度n。
3. 计算结果：每个元素的最长符合条件子数组长度为right[i]-left[i]-1，累加所有长度得到最终结果。

内容的提问来源于stack exchange，提问作者Learner