和至少为k的最短子数组 python_LeetCode 862. 和至少为 K 的最短子数组

题目描述

返回 A 的最短的非空连续子数组的 长度，该子数组的和至少为 K 。

如果没有和至少为 K 的非空子数组，返回 -1 。

样例

输入：A = [1], K = 1

输出：1

输入：A = [1,2], K = 4

输出：-1

输入：A = [2,-1,2], K = 3

输出：3

注意

1 <= A.length <= 50000

-10 ^ 5 <= A[i] <= 10 ^ 5

1 <= K <= 10 ^ 9

算法1

(树状数组) $O(n \log n)$

构造前缀和数组 $s(i)$，对于每个 $i$，找到下标最大的 $j$, $(j

将所有 $0, s(i), K, s(i) + K$ 进行离散化，离散化到 $[1, 2n + 2]$，然后使用权值树状数组寻找最大的 $j$。

具体地，在对应 $s$ 值的树状数组位置上更新最大值。查询时从树状数组中寻找一个前缀最大值，然后再更新树状数组即可。

时间复杂度

每次更新和查询的时间复杂度为 $O(\log n)$，故总时间复杂度为 $O(n \log n)$。

C++ 代码

class Solution {

public:

void update(vector &f, int x, int y) {

for (; x < f.size(); x += x & -x)

f[x] = max(f[x], y);

}

int query(vector &f, int x) {

int t = -1;

for (; x; x -= x & -x)

t = max(f[x], t);

return t;

}

int shortestSubarray(vector& A, int K) {

int n = A.size();

vector s(2 * n + 2, 0), d(2 * n + 2, 0);

vector f(2 * n + 3, -1);

for (int i = 1; i <= n; i++) {

s[i] = s[i - 1] + A[i - 1];

d[i] = s[i];

}

for (int i = n + 1; i <= 2 * n + 1; i++) {

s[i] = s[i - n - 1] + K;

d[i] = s[i];

}

sort(d.begin(), d.end());

for (int i = 0; i <= 2 * n + 1; i++)

s[i] = lower_bound(d.begin(), d.end(), s[i]) - d.begin() + 1;

update(f, s[n + 1], 0);

int ans = n + 1;

for (int i = 1; i <= n; i++) {

int t = query(f, s[i]);

if (t != -1)

ans = min(ans, i - t);

update(f, s[i + n + 1], i);

}

if (ans == n + 1)

ans = -1;

return ans;

}

};

算法2

(单调队列) $O(n)$

求出前缀和数组 $s$。

对于某个 $s(i)$，我们期望有一个单调的队列，从 $j = q.front()$ 开始，在满足 $s(j) + K \le s(i)$ 的情况下一直向后移动 $j$，直到条件不满足，此时 $s(i)$ 所能更新的答案就是 $i - j$。

现在考察这个队列的性质，为什么队列一定是单调递增的？这是因为如果出现了递减的情况，后一个元素一定比前一个元素更优(值更小更容易满足 $s(j) + K \le s(i)$，且 $j$ 更靠近 $i$)所以递减时，前一个元素是没有意义的。

根据这个性质，每次 $s(i)$ 入队时，维护队列的单调性。

在每次 $j$ 从队头遍历的过程中，是不是被遍历过元素可以丢弃呢？可以。这是因为被遍历的元素是不会再产生更优的元素了，下一次的 $i$ 一定更大，与这个 $j$ 无论如何都不会凑出更小的答案。

时间复杂度

每个元素最多进队一次，出队一次，故时间复杂度为 $O(n)$。

空间复杂度

需要 $O(n)$ 的额外空间存储前缀和数组和单调队列。

C++ 代码

class Solution {

public:

int shortestSubarray(vector& A, int K) {

int n = A.size();

vector s(n + 1, 0);

deque q;

for (int i = 1; i <= n; i++)

s[i] = s[i - 1] + A[i - 1];

q.push_back(0);

int ans = n + 1;

for (int i = 1; i <= n; i++) {

while (!q.empty() && s[q.front()] + K <= s[i]) {

ans = min(ans, i - q.front());

q.pop_front();

}

while (!q.empty() && s[q.back()] >= s[i])

q.pop_back();

q.push_back(i);

}

if (ans == n + 1)

ans = -1;

return ans;

}

};