Maximum Subarray Sum After at Most K Swaps
Maximum Subarray Sum After at Most K Swaps
You are given an integer array nums and an integer k.
You are allowed to perform at most k swap operations on the array.
In one swap operation, you may choose any two indices i and j and swap nums[i] and nums[j].
Return an integer denoting the maximum possible subarray sum after performing the swaps.
Example 1:
Input: nums = [1,-1,0,2], k = 1
Output: 3
Explanation:
- We can swap on indices 1 and 3, resulting in the array
[1, 2, 0, -1]. - The subarray
[1, 2]has a sum of 3, which is the maximum possible subarray sum after at mostk = 1 swap.
Example 2:
Input: nums = [4,3,2,4], k = 2
Output: 13
Explanation:
The maximum possible subarray sum after at most k = 2 swaps is the sum of the entire array, which is 13.
Example 3:
Input: nums = [-1,-2], k = 0
Output: -1
Explanation:
k = 0swaps are allowed.- The possible subarrays are
[-1],[-2], and[-1, -2], with sums -1, -2, and -3 respectively. - Among these sums, the maximum is -1.
Constraints:
1 <= nums.length <= 1500-105 <= nums[i] <= 1050 <= k <= nums.length
解题思路
首先，任何被交换进子数组的元素最终必然会留在选定的子数组中，否则交换操作没有意义。为了找到能使子数组和最大的方案，我们需要确定该选择哪个子数组。注意到数组长度 $n$ 最大只有 $1500$，这允许我们可以通过 $O(n^2)$ 的复杂度暴力枚举所有可能的子数组，从而对每个子数组计算在至多 $k$ 次交换下所能获得的最大子数组和。
在具体实现上，我们可以先枚举子数组的左端点 $l$，随后逐步向右扩展右端点 $r$。对于任意固定的子数组区间 $[l,r]$，交换子数组内部的两个元素显然无法增加该区间的元素总和，因此有效的交换只会发生在子数组内部与外部元素之间。为了更清晰地描述这一过程，记 $S_{[l,r]}$ 为子数组内元素构成的非递减有序可重集合，其中 $S_{[l,r]}^{(i)}$ 表示集合中第 $i$ 小的元素；同时记 $S_{\overline{[l,r]}}$ 为子数组外其余元素构成的非递增有序可重集合，其中 $S_{\overline{[l,r]}}^{(i)}$ 表示集合中第 $i$ 大的元素。基于贪心思想，为了使子数组 $[l,r]$ 的和最大，我们应选择 $S_{[l,r]}$ 中最小的 $c$ 个数与 $S_{\overline{[l,r]}}$ 中最大的 $c$ 个数进行交换。这里的 $c$ 需满足约束：
- 交换次数不能超过限制，即 $c \le k$；
- 参与交换的元素数量不能超过各自集合的大小，即 $c \le \min\{|S_{[l,r]}|, |S_{\overline{[l,r]}}|\}$；
- 每次交换必须带来正收益，即对于任意 $1 \le i \le c$，都要满足 $S_{[l,r]}^{(i)} < S_{\overline{[l,r]}}^{(i)}$。
对于给定的子数组，我们需要快速求出满足上述条件的最大交换次数 $c$。定义差值 $d_i = S_{[l,r]}^{(i)} - S_{\overline{[l,r]}}^{(i)}$，随着 $i$ 的增大，$S_{[l,r]}^{(i)}$ 逐渐变大而 $S_{\overline{[l,r]}}^{(i)}$ 逐渐变小，因此 $d_i$ 具有单调不减的性质。这意味着存在一个分界点 $c$，使得当 $i \le c$ 时 $d_i < 0$，而当 $i \ge c+1$ 时 $d_i \ge 0$。利用这种二段性，我们可以通过二分找到最大的 $c$。同时为了支持在枚举过程中动态维护集合并快速查询第 $k$ 小元素及前 $k$ 小元素之和，我们可以借助平衡树或值域树状数组等数据结构，将这些操作的时间复杂度均控制在 $O(\log n)$ 级别。
上述思路的整体时间复杂度为 $O(n^2 \log^2 n)$，其中枚举子数组的复杂度为 $O(n^2)$，而针对每个子数组进行二分并结合数据结构查询第 $k$ 小元素的复杂度为 $O(\log^2 n)$。实际上，我们可以进一步优化掉二分，将整体复杂度降至 $O(n^2 \log n)$。当固定左端点 $l$ 并向右扩展右端点 $r$ 时，每次扩展实质上是将一个元素从外部集合移入内部集合。可以证明，这一操作使得最优交换次数 $c$ 的变化幅度最多为 $1$。因此，在每次向右扩展时，我们只需检查将 $c$ 增加 $1$ 或减少 $1$ 是否依然满足条件，并据此更新 $c$ 即可，无需重新进行二分。对于博主来说这个结论很难从直觉上注意到，所以下面给出严谨的证明。
对于固定的左端点 $l$，当右端点从 $r$ 扩展到 $r+1$ 时，内部集合加入一个元素，外部集合删除一个元素。对于内部集合，新集合的第 $i$ 小元素 $S_{[l,r+1]}^{(i)}$ 只能来自原集合的第 $i$ 小、第 $i-1$ 小或新加入的元素，故必然满足 $S_{[l,r]}^{(i-1)} \le S_{[l,r+1]}^{(i)} \le S_{[l,r]}^{(i)}$。对于外部集合，新集合的第 $i$ 大元素 $S_{\overline{[l,r+1]}}^{(i)}$ 同理必然满足 $S_{\overline{[l,r]}}^{(i)} \le S_{\overline{[l,r+1]}}^{(i)} \le S_{\overline{[l,r]}}^{(i-1)}$。
将这两个不等式相减，可得 $S_{[l,r]}^{(i-1)} - S_{\overline{[l,r]}}^{(i)} \le S_{[l,r+1]}^{(i)} - S_{\overline{[l,r+1]}}^{(i)} \le S_{[l,r]}^{(i)} - S_{\overline{[l,r]}}^{(i-1)}$。又因为 $S_{[l,r]}^{(i-1)} - S_{\overline{[l,r]}}^{(i)} \ge S_{[l,r]}^{(i-1)} - S_{\overline{[l,r]}}^{(i-1)}$ 且 $S_{[l,r]}^{(i)} - S_{\overline{[l,r]}}^{(i-1)} \le S_{[l,r]}^{(i)} - S_{\overline{[l,r]}}^{(i)}$，最终得到 $S_{[l,r]}^{(i-1)} - S_{\overline{[l,r]}}^{(i-1)} \le S_{[l,r+1]}^{(i)} - S_{\overline{[l,r+1]}}^{(i)} \le S_{[l,r]}^{(i)} - S_{\overline{[l,r]}}^{(i)}$。记 $d_i = S_{[l,r]}^{(i)} - S_{\overline{[l,r]}}^{(i)}$，$d'_i = S_{[l,r+1]}^{(i)} - S_{\overline{[l,r+1]}}^{(i)}$，则上述不等式等价于 $d_{i-1} \le d'_i \le d_i$。
假设在扩展到 $r+1$ 之前，最优交换次数为 $c$，即序列 $d$ 的前 $c$ 项为负数，第 $c+1$ 项起为非负数，满足 $d_c < 0$ 且 $d_{c+1} \ge 0$。在扩展到 $r+1$ 后，根据推论 $d'_{c+2} \ge d_{c+1} \ge 0$，这意味着新序列第 $c+2$ 项及之后的所有项均非负，负数最多只有 $c+1$ 个，因此 $c$ 最多增加 $1$；同时，由于 $d'_{c-1} \le d_{c-1} \le d_c < 0$，新序列第 $c-1$ 项及之前的项均为负数，负数至少有 $c-1$ 个，因此 $c$ 最多减少 $1$。
所以，当右端点扩展时，最优交换次数 $c$ 的变化幅度不超过 $1$。得证。
AC 代码如下，时间复杂度为 $O(n^2\log{n})$：
const int N = 1505;
int xs[N], sz;
struct Fenwick {
long long tr1[N], tr2[N];
Fenwick() {
memset(tr1, 0, sz + 1 << 3);
memset(tr2, 0, sz + 1 << 3);
}
int lowbit(int x) {
return x & -x;
}
void add(int x, int c, int v) {
for (int i = x; i <= sz; i += lowbit(i)) {
tr1[i] += c;
tr2[i] += v;
}
}
int kth(int k) {
int ret = 0;
for (int i = __lg(sz); i >= 0; i--) {
int t = ret | 1 << i;
if (t <= sz && tr1[t] < k) {
ret = t;
k -= tr1[t];
}
}
return ret + 1;
}
long long query(int k) { // 求前k小元素之和。采用与kth相同的二分方式，处理第k小有多个重复值的情况
if (!k) return 0;
long long x = 0, ret = 0;
for (int i = __lg(sz); i >= 0; i--) {
int t = x | 1 << i;
if (t <= sz && tr1[t] < k) {
x = t;
ret += tr2[t];
k -= tr1[t];
}
}
ret += 1ll * xs[x + 1] * k; // 此时k为剩余需累加的元素个数，xs[x+1]为第k小的值
return ret;
}
};
class Solution {
public:
long long maxSum(vector<int>& nums, int k) {
int n = nums.size();
sz = 0;
for (auto &x : nums) {
xs[++sz] = x;
}
sort(xs + 1, xs + sz + 1);
sz = unique(xs + 1, xs + sz + 1) - xs - 1;
for (auto &x : nums) {
x = lower_bound(xs + 1, xs + sz + 1, x) - xs;
}
long long ret = -1e18;
for (int i = 0; i < n; i++) {
Fenwick tr1, tr2;
for (auto &x : nums) {
tr2.add(x, 1, xs[x]);
}
long long s = 0;
for (int j = i, c = 0; j < n; j++) {
tr1.add(nums[j], 1, xs[nums[j]]);
tr2.add(nums[j], -1, -xs[nums[j]]);
s += xs[nums[j]];
int t = j - i + 1, flag = 0;
if (c < k && c < t && c < n - t && tr1.kth(c + 1) < tr2.kth(n - t - c)) c++, flag = 1;
if (!flag && c && tr1.kth(c) >= tr2.kth(n - t - c + 1)) c--;
ret = max(ret, s + tr2.query(n - t) - tr2.query(n - t - c) - tr1.query(c));
}
}
return ret;
}
};
参考资料
至多 K 次交换后最大子数组和【力扣周赛 506】：https://www.bilibili.com/video/BV1ptJw6hENZ/
本文来自博客园，作者：onlyblues，转载请注明原文链接：https://www.cnblogs.com/onlyblues/p/20580953
浙公网安备 33010602011771号