Medium

题目描述

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

对于每个索引 i（从 0 到 n-1），执行以下操作：

  • 找到所有满足 nums1[j] < nums1[i] 的索引 j
  • 在这些索引中最多选择 k 个 nums2[j] 的值来最大化总和

返回一个大小为 n 的数组 answer，其中 answer[i] 表示对应索引 i 的结果。

示例 1：

输入: nums1 = [4,2,1,5,3], nums2 = [10,20,30,40,50], k = 2
输出: [80,30,0,80,50]
解释:
- 对于 i = 0: 从索引 [1,2,4] 中选择 2 个最大值（nums1[j] < nums1[0]），结果为 50 + 30 = 80
- 对于 i = 1: 从索引 [2] 中选择 2 个最大值（nums1[j] < nums1[1]），结果为 30
- 对于 i = 2: 没有索引满足 nums1[j] < nums1[2]，结果为 0
- 对于 i = 3: 从索引 [0,1,2,4] 中选择 2 个最大值（nums1[j] < nums1[3]），结果为 50 + 30 = 80
- 对于 i = 4: 从索引 [1,2] 中选择 2 个最大值（nums1[j] < nums1[4]），结果为 30 + 20 = 50

示例 2：

输入: nums1 = [2,2,2,2], nums2 = [3,1,2,3], k = 1
输出: [0,0,0,0]
解释: 由于 nums1 中所有元素都相等，没有索引满足条件 nums1[j] < nums1[i]，所以所有位置的结果都是 0

约束条件：

  • n == nums1.length == nums2.length
  • 1 <= n <= 10^5
  • 1 <= nums1[i], nums2[i] <= 10^6
  • 1 <= k <= n

解题思路

这道题的关键在于高效地处理每个位置的查询。暴力解法会超时，需要优化思路。

核心思想：排序 + 优先队列

  1. 预处理排序：将数组元素按照 nums1 的值进行排序，同时保存原始索引。这样可以按照 nums1 的值从小到大处理。

  2. 使用最大堆维护候选值：对于当前处理的元素，所有之前处理过的元素的 nums1 值都小于当前值，因此它们的 nums2 值都是候选值。使用最大堆来维护这些候选值中的前 k 个最大值。

  3. 动态更新堆：当处理到新元素时，将其 nums2 值加入堆中。如果堆的大小超过 k，则移除最小值。堆中所有元素的和就是当前位置的答案。

  4. 处理相同值：对于 nums1 值相同的元素，需要特别处理，因为它们之间不能互相选择。

时间复杂度优化：通过排序预处理，每个元素只需要O(log k)的时间来维护堆，总体复杂度为O(n log n + n log k)。

推荐解法：排序 + 最大堆，既高效又容易理解。

代码实现

class Solution {
public:
    vector<long long> findMaxSum(vector<int>& nums1, vector<int>& nums2, int k) {
        int n = nums1.size();
        vector<pair<int, int>> sorted_pairs;
        
        // 创建 (nums1[i], i) 对并排序
        for (int i = 0; i < n; i++) {
            sorted_pairs.push_back({nums1[i], i});
        }
        sort(sorted_pairs.begin(), sorted_pairs.end());
        
        vector<long long> answer(n, 0);
        priority_queue<int> max_heap; // 最大堆
        long long current_sum = 0;
        
        int i = 0;
        while (i < n) {
            int current_val = sorted_pairs[i].first;
            
            // 处理所有与当前值相同的元素
            vector<int> same_group_indices;
            while (i < n && sorted_pairs[i].first == current_val) {
                same_group_indices.push_back(sorted_pairs[i].second);
                i++;
            }
            
            // 为同组的每个元素计算答案
            for (int idx : same_group_indices) {
                answer[idx] = current_sum;
            }
            
            // 将同组元素的 nums2 值添加到堆中
            for (int idx : same_group_indices) {
                max_heap.push(nums2[idx]);
                current_sum += nums2[idx];
                
                // 如果堆大小超过 k，移除最小元素
                if (max_heap.size() > k) {
                    current_sum -= max_heap.top();
                    max_heap.pop();
                }
            }
        }
        
        return answer;
    }
};

class Solution:
    def findMaxSum(self, nums1: List[int], nums2: List[int], k: int) -> List[int]:
        import heapq
        
        n = len(nums1)
        # 创建 (nums1[i], i) 对并排序
        sorted_pairs = sorted((nums1[i], i) for i in range(n))
        
        answer = [0] * n
        min_heap = []  # 用最小堆模拟最大堆（存储负值）
        current_sum = 0
        
        i = 0
        while i < n:
            current_val = sorted_pairs[i][0]
            
            # 处理所有与当前值相同的元素
            same_group_indices = []
            while i < n and sorted_pairs[i][0] == current_val:
                same_group_indices.append(sorted_pairs[i][1])
                i += 1
            
            # 为同组的每个元素计算答案
            for idx in same_group_indices:
                answer[idx] = current_sum
            
            # 将同组元素的 nums2 值添加到堆中
            for idx in same_group_indices:
                heapq.heappush(min_heap, -nums2[idx])  # 存储负值实现最大堆
                current_sum += nums2[idx]
                
                # 如果堆大小超过 k，移除最小元素
                if len(min_heap) > k:
                    current_sum += heapq.heappop(min_heap)  # 加上负值等于减去正值
        
        return answer

public class Solution {
    public long[] FindMaxSum(int[] nums1, int[] nums2, int k) {
        int n = nums1.Length;
        var sortedPairs = new List<(int val, int idx)>();
        
        // 创建 (nums1[i], i) 对并排序
        for (int i = 0; i < n; i++) {
            sortedPairs.Add((nums1[i], i));
        }
        sortedPairs.Sort();
        
        long[] answer = new long[n];
        var maxHeap = new PriorityQueue<int, int>(Comparer<int>.Create((a, b) => b.CompareTo(a)));
        long currentSum = 0;
        
        int i = 0;
        while (i < n) {
            int currentVal = sortedPairs[i].val;
            
            // 处理所有与当前值相同的元素
            var sameGroupIndices = new List<int>();
            while (i < n && sortedPairs[i].val == currentVal) {
                sameGroupIndices.Add(sortedPairs[i].idx);
                i++;
            }
            
            // 为同组的每个元素计算答案
            foreach (int idx in sameGroupIndices) {
                answer[idx] = currentSum;
            }
            
            // 将同组元素的 nums2 值添加到堆中
            foreach (int idx in sameGroupIndices) {
                maxHeap.Enqueue(nums2[idx], nums2[idx]);
                currentSum += nums2[idx];
                
                // 如果堆大小超过 k，移除最小元素
                if (maxHeap.Count > k) {
                    currentSum -= maxHeap.Dequeue();
                }
            }
        }
        
        return answer;
    }
}

var findMaxSum = function(nums1, nums2, k) {
    const n = nums1.length;
    const answer = new Array(n);
    
    for (let i = 0; i < n; i++) {
        const candidates = [];
        
        for (let j = 0; j < n; j++) {
            if (nums1[j] < nums1[i]) {
                candidates.push(nums2[j]);
            }
        }
        
        candidates.sort((a, b) => b - a);
        
        let sum = 0;
        for (let idx = 0; idx < Math.min(k, candidates.length); idx++) {
            sum += candidates[idx];
        }
        
        answer[i] = sum;
    }
    
    return answer;
};
在 LeetCode 上运行： 力扣中文版 LeetCode 国际版

复杂度分析

复杂度类型复杂度说明
时间复杂度O(n log n + n log k)排序需要 O(n log n)，每个元素的堆操作需要 O(log k)
空间复杂度O(n + k)存储排序数组需要 O(n)，堆的大小最多为 k