如何求和第K大(小)的子序列

问题来源

此题是上场$leetcode$的最后一题
https://leetcode.cn/problems/find-the-k-sum-of-an-array/

另一个问题

• 上面那道题看完感觉很经典，但是想了好久好像做过类似的，$tag$是优先队列，那么先看一下我曾经做过的那个题
  https://www.luogu.com.cn/problem/P1631
• 让你从两个序列中各取一个数加和，这会得到$N^2$个数，找最小的$N$个数
• 这个题目的思路是我们可以先把所有的$a[i]+b[1]$放到一个小顶堆里面，这样的话堆顶就是最小的数，同时我们记录这个数对应的$a$和$b$数组的下标，如果当前达到最小，那么由于数组都是有序的，所以下一个最小的只需要把$b$下标+1即可，容易理解

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;

struct st{
  int a, b;
  ll sum;
  bool operator < (const st &B)const{
    return sum > B.sum;
  }
};
int main(){
  ios::sync_with_stdio(false);
  cin.tie(0);
  cout.tie(0);
  int n;
  cin >> n;
  vector<ll> a(n + 1), b(n + 1);
  for(int i=1;i<=n;i++) cin >> a[i];
  for(int i=1;i<=n;i++) cin >> b[i];
  priority_queue<st> q;
  for(int i=1;i<=n;i++){
    q.push(st{i, 1, a[i] + b[1]});
  }
  for(int i=1;i<=n;i++){
    auto u = q.top();
    q.pop();
    cout << u.sum << " \n"[i == n];
    q.push(st{u.a, u.b + 1, b[u.b + 1] + a[u.a]}); 
  }
  return 0;
}

• 这两个题有相似之处，下面进行分析

k大和

• 考虑求出总和减去第$k$小就是第$k$大

无负数

• 假设整个数组都是非负数，如何考虑？在这种情况下，如果整个数组升序排列，那么最小的子序列显然是$a[0]$，那么问题来了，第二小的是谁？我们有两种选择:$a[1],a[0]+a[1]$，其余不可能成为第二小，所以我们可以仿照上面的解法，在一个小根堆里面放一个二元组，$(s[i],i)$表示当前子序列和为$s[i]$，最后一个元素下标为$i$，按照上面的两种方式进行进堆出堆，可以发现这样对于子序列的考虑是完备的，如果不考虑空集，弹出的第$k$个就是第$k$小
• 形式化描述就是当前堆顶元素为$(s[i],i)$，接下来入堆的元素就是$(s[i]+a[i+1],i+1)$或者$(s[i]-a[i]+a[i+1],i+1)$，比较它们两个大小或者直接全入堆，最多操作$2k$次

带负数

• 加上负数主要的困难在于多个负数累加起来会变得更小，不满足全局单调性，给入堆和出堆的方案选择造成了困难，那么我们可以换一种角度考虑，如何把问题转化回我们熟悉的无负数问题。首先负数化正，然后记录一下负数和，求完无负数的问题，每个值加上负数和即为答案，这个怎么理解呢？可以这样想，对于每一个子序列，我加上任意正数或负数影响都是一样的，那么我通过把负数变正数的影响去掉，就可以得到变号之前的答案了

typedef pair<long long, int> pli;
class Solution {
public:
    long long kSum(vector<int>& a, int k) {
        int n = (int)a.size();
        long long ans = 0;
        long long sum = 0;
        long long neg = 0;
        for(int i=0;i<n;i++){
            sum += a[i];            
            if(a[i] < 0){
                neg += a[i];
                a[i] = -a[i];
            }
        }sort(a.begin(), a.end(), less<int>());
        priority_queue<pli, vector<pli>, greater<pli> > q;// 小根堆
        q.push(make_pair(a[0], 0));
        for(int i=0;i<k-1;i++){
            auto u = q.top();
            q.pop();
            ans = u.first;
            if(u.second == n - 1) continue;
            q.push(make_pair(u.first + a[u.second + 1], u.second + 1));
            q.push(make_pair(u.first - a[u.second] + a[u.second + 1], u.second + 1));
        }
        return sum - (neg + ans);
    }
};