求数组第k小/大的序列和

Posted 2022-08-23 Harris-H

求数组第$k$小/大的序列和

弱化版：$a_i$ 非负，求第$k$小。

最小的子序列和是空集$0$。

算法流程：

使用优先队列，用二元组$(sum,i)$ 表示当前下标$i$结尾的子序列和$sum$。

初始加入$(0,0)$。

每次取出队首。

然后加入$(sum+a_{i+1},i+1),(sum+a_{i+1}-a_i,i+1)$。

第$k$次即为第$k$小的序列和。

时间复杂度：$O(n\log n+k\log k)$

证明

我们需要证明这样做得到的子序列是不重复、不遗漏、按顺序。

---

强化版1：$a_i$ 为整数，求第$k$大。

显然可以把$a_i$分成两部分，负数和非负数。

显然第$1$大的和就是$sum(a_i),a_i\ge 0$。

接下来要么减去一个非负数，要么加上一个负数。

这两种情况可以等价为减上一个非负数。然后我们对这些非负数从小到大排序。

然后就转化弱化版的题目的做法，使用优先队列即可。

代码

class Solution 
public:
    long long kSum(vector<int> &nums, int k) 
        long sum = 0L;
        for (int &x : nums)
            if (x >= 0) sum += x;
            else x = -x;
        sort(nums.begin(), nums.end());
        priority_queue<pair<long, int>> pq;
        pq.emplace(sum, 0);
        while (--k) 
            auto[sum, i] = pq.top();
            pq.pop();
            if (i < nums.size()) 
                pq.emplace(sum - nums[i], i + 1); // 保留 nums[i-1]
                if (i) pq.emplace(sum - nums[i] + nums[i - 1], i + 1); // 不保留 nums[i-1]
            
        
        return pq.top().first;
    
;

强化版2：$a_i$ 为整数，求第$k$小。

类似地。最小的就是负数之和。

接下可以等价为加一个非负数，对这些非负数从小到大排序。

然后转弱化版解法。

代码略。

---