给定列表中查找满足阈值限制的最长连续元素范围算法优化

问题定义 #

查找满足指定条件的最长连续元素。
现有数字列表A、数字列表B，以及阈值Limit。
任务为找出A中最长的k个连续元素，满足如下条件：

Max(B[i],B[i+1],...B[i+k]) + Sum(A[i], A[i+1], ..., A[i+k]) * k ≤ Limit

示例说明 #

A = [2,1,3,4,5]
B = [3,6,1,3,4]
Limit = 25

取2个连续元素：
A中对应和最高的元素为4、5，对应B段的最大值为Max(3,4) = 4。
计算值为4 + (4+5) * 2 = 22 ≤ 25，因此2个连续元素符合要求。

取3个连续元素：
取A前3个元素2、1、3，对应B段的最大值为Max(3,6,1) = 6。
计算值为6 + (2+1+3) * 3 = 24 ≤ 25，因此3个连续元素符合要求。

取4个连续元素：
取A前4个元素2、1、3、4，对应B段的最大值为Max(3,6,1,3) = 6。
计算值为6 + (2+1+3+4) * 4 = 46 > 25，因此4个连续元素不符合要求。

所以该输入的正确答案为3。

约束条件 #

• n（A的长度）最高为10⁵
• A元素最大值为10¹⁴
• B元素最大值为10⁹
• Limit最大值为10¹⁴

初始代码问题 #

最初的暴力实现时间复杂度为O(n³)，三层循环遍历所有长度、所有起始位置、再计算区间和与最大值，n超过1e3就会出现明显超时，完全无法适配1e5的输入规模，代码如下：

public int getMax(List<Integer> A, List<Integer> B, long limit) {
    int result = 0;
    int n = A.size();
    for(int len=1; len<=n; len++) {
        for(int i=0; i<=n-len; i++) {
            int j=i+len-1;
            int max = B.get(i);
            long total = 0;
            for(int k=i; k<=j; k++) {
                total += A.get(k);
                max = Math.max(max, B.get(k));
            }
            total = max + total * len;
            if(total < limit) {
                result = len;
                break;
            }
        }
    }
    return result;
}

滑动窗口修改版本错误原因 #

修改后的滑动窗口代码返回结果错误，同时仍然没有解决超时问题，核心错误点有两个：

1. 长度计算逻辑不一致：计算校验值的时候用的是to - from作为窗口长度，但更新最大长度的时候用的是to - from +1，导致校验的是长度为k的窗口，却误认为是k+1的窗口，这是示例返回4的直接原因。
2. 每次窗口变动都重新遍历计算区间和与最大值，时间复杂度还是O(n²)，依然会在1e5输入规模下超时。
3. 普通滑动窗口没有维护区间最大值的能力，左边界收缩时无法快速获取新窗口的B最大值，必须借助单调队列实现。

修改后的错误代码如下：

public int getMax(List<Integer> A, List<Integer> B, long limit) {
    int from = 0, to = 0, max = -1;
    int n = A.size();
    for (; from < n;) {
        int total = 0;
        int m = B.get(from); // updated here
        for (int i = from; i < to; i++) {
            total += A.get(i); // updated here
            m = Math.max(m, B.get(i)); // updated here
        }

        total = m + total * (to - from); // updated here

        if (total <= limit && to - from + 1 > max) {
            max = to - from + 1;
        }
        if (total < limit && to < n) { // below target, extend window
            to++;
        } else { // otherwise contract window
            from++;
        }
        if (from > to) {
            to = from;
        }
    }
    return max;
}

正确实现方案 #

采用前缀和+单调队列维护滑动窗口的方案，时间复杂度为O(n)，可以完美适配1e5的输入规模：

1. 预先计算A的前缀和数组，实现O(1)获取任意区间的A元素和
2. 用单调递减双端队列维护当前滑动窗口内B的最大值，队首始终是当前窗口B的最大值下标，新增元素和收缩左边界时都可以O(1)维护队列
3. 采用变长滑动窗口遍历，右边界不断右移，当窗口不满足条件时收缩左边界，每次满足条件时更新最大长度

正确代码实现 #

import java.util.Deque;
import java.util.LinkedList;
import java.util.List;

public int getMax(List<Integer> A, List<Integer> B, long limit) {
    int n = A.size();
    if (n == 0) return 0;
    // 前缀和数组，preSum[i] = A[0]+A[1]+...+A[i-1]
    long[] preSum = new long[n + 1];
    for (int i = 0; i < n; i++) {
        preSum[i + 1] = preSum[i] + A.get(i);
    }
    Deque<Integer> deque = new LinkedList<>();
    int left = 0;
    int maxLen = 0;
    for (int right = 0; right < n; right++) {
        // 维护单调递减队列，新增元素B[right]
        while (!deque.isEmpty() && B.get(deque.peekLast()) <= B.get(right)) {
            deque.pollLast();
        }
        deque.offerLast(right);
        // 收缩左边界直到窗口满足条件
        while (left <= right) {
            // 当前窗口最大值
            int currentMaxB = B.get(deque.peekFirst());
            int len = right - left + 1;
            long sumA = preSum[right + 1] - preSum[left];
            long total = currentMaxB + sumA * len;
            if (total <= limit) {
                break;
            }
            // 不满足条件，左边界右移
            if (deque.peekFirst() == left) {
                deque.pollFirst();
            }
            left++;
        }
        // 更新最大长度
        maxLen = Math.max(maxLen, right - left + 1);
    }
    return maxLen;
}

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