Skip to content

第 307 场周赛

T4 找出数组的第 K 大和

Link

给你一个整数数组 nums 和一个 正整数 k 。你可以选择数组的任一子序列 并且对其全部元素求和。

数组的 第 k 大和 定义为:可以获得的第 k 个 最大 子序列和(子序列和允许出现重复)

返回数组的 第 k 大和 。

子序列是一个可以由其他数组删除某些或不删除元素排生而来的数组,且派生过程不改变剩余元素的顺序。

注意:空子序列的和视作 0 。

\(1\le n \le 10^5, -10^9 \le nums[i] \le 10^9, k \le min(2000, 2^n)\)


Step 1:

问题中包含了负数,直观上如果只包含正数问题会简化很多,所以第一步思考如何把负数影响消去。我们假设负数都必须选,令 sum 为所有负数的和,然后把负数都变为其相反数。这样,现在这个全部都是正数的数组中的某个子数组和,加上 sum 就可以与原数组的子数组和一一对应。(选上某个负数的相反数,对应于原数组中不选择这个负数)

Step 2:

问题要求解第 \(k\) 大和,对于一个数字都是正数的序列来说,最大和无非是所有数字都选择的情况,然后次大是从中去掉某些值。情况很复杂,但注意到 k 最多只有 2000,提示我们可以一一遍历这些值。

Step 3:

求解第 \(k\) 大,其实等同于求解第 \(k\) 小(允许不选的情况下),只需要用总和减去就可以相互转换。

至此,问题简化为在一个全是整数的数组中,求解第 k 小的子数组和。并且允许子数组为空。

不妨将数组从小到大排序,什么数字都不选就是最小的答案 0,紧接着次小就是只选择 nums[0] 。但再往后就复杂了起来,在 nums[0] 的基础上,往后延伸变为 nums[0] + nums[1],或者只是 nums[1]。这时就产生了两个以 1 结尾的答案。

根据这条规律,对于以 \(i\) 结尾的子数组和 \(s\),可以产生 \(s + nums[i + 1]\)\(s + nums[i + 1] - nums[i]\) 这两个以 \(i+1\) 结尾的子数组和。并且这个拓展方向是递增的,后产生的永远比先产生的数值大。为了保证整体扩展的有序性,我们用最小堆来做这件事情。

另外每次规模都会乘以 2,这保证了所有方案可以不重不漏的进行枚举。

总体复杂度为 \(O(n\log n +k\log k)\)

class Solution {
public:
    using ll = long long;
    using PLI = pair<ll, int>;
    long long kSum(vector<int>& nums, int k) {
        ll sum = 0, res = 0;
        for(auto &x : nums) {
            if(x < 0) {
                sum += x;
                x = -x;
            }
        }
        sort(nums.begin(), nums.end());

        priority_queue<PLI, vector<PLI>, greater<PLI>> q;
        q.push({0, -1});
        while(k) {
            auto [s, i] = q.top(); q.pop(); k--;
            res = s;
            if(i + 1 < nums.size()) {
                q.push({s + nums[i + 1], i + 1});
                if(i >= 0) q.push({s + nums[i + 1] - nums[i], i + 1}); // 注意 i = -1 时需要跳过
            }
        }
        return accumulate(nums.begin(), nums.end(), 0ll) - res + sum;
    }
};
本文访问