3478. Choose K Elements With Maximum Sum

You are given two integer arrays, nums1 and nums2, both of length n, along with a positive integer k.

For each index i from 0 to n - 1, perform the following:

• Find all indices j where nums1[j] is less than nums1[i].
• Choose at most k values of nums2[j] at these indices to maximize the total sum.

Return an array answer of size n, where answer[i] represents the result for the corresponding index i.

c++

class Solution {
public:
    vector<long long> findMaxSum(vector<int>& nums1, vector<int>& nums2, int k) {
        vector<array<long long,3>> S;
        long long n = nums1.size(), sum = 0;
        priority_queue<long long,vector<long long>,greater<>> q;
        for(int i = 0; i < n; i++) S.push_back({nums1[i], nums2[i], i});
        sort(rbegin(S), rend(S));
        vector<long long> res(n);
        while(S.size()) {
            auto x = S.back()[0];
            long long ans = sum;
            while(S.size() and S.back()[0] == x) {
                auto [_,now,idx] = S.back(); S.pop_back();
                res[idx] = ans;
                sum += now;
                q.push(now);
            }
            while(q.size() > k) {
                sum -= q.top(); q.pop();
            }
        }
        return res;
    }
};