Leetcode 3478

题目要求

给你两个整数数组，nums1 和 nums2，长度均为 n，以及一个正整数 k 。

对从 0 到 n - 1 每个下标 i ，执行下述操作：

• 找出所有满足 nums1[j] 小于 nums1[i] 的下标 j 。

• 从这些下标对应的 nums2[j] 中选出 至多 k 个，并 最大化 这些值的总和作为结果。

返回一个长度为 n 的数组 answer ，其中 answer[i] 表示对应下标 i 的结果。

核心算法和原理

• 数据离散化的非常简单处理：用vector<pairs<int, int>>装 (nums[i], i)，sort即可

代码

class Solution {
public:
    vector<long long> findMaxSum(vector<int>& nums1, vector<int>& nums2, int k) {
        int n = nums1.size();
        vector<long long> ans(n);
        typedef pair<int, int> pii;
        vector<pii> sn1;
        for (int i = 0; i < n; i++) {
            sn1.emplace_back(nums1[i], i);
        }
        
        ranges::sort(sn1);
        

        priority_queue<int, vector<int>, greater<int>> pq;

        long long sum = 0;
        for (int i = 0; i < n; i++) {
            // update ans[p.second]
            pii p = sn1[i];
            ans[p.second] = (i > 0 && sn1[i].first == sn1[i - 1].first) ? ans[sn1[i - 1].second] : sum;
            int n2 = nums2[p.second];
            if (pq.size() < k) {
                sum += n2;
                pq.push(n2);
            } else if (pq.top() < n2) {
                sum -= pq.top();
                sum += n2;
                pq.pop();
                pq.push(n2);
            }
        }
        return ans;
    }
};