跳转至

3763. 带阈值约束的最大总和 🔒

题目描述

给定两个整数数组 nums 和 threshold，长度都是 n。

从 step = 1 开始，重复执行下面操作：

  • 找到一个 未使用 的下标 i 使得 threshold[i] <= step。
    • 如果没有这样的下标存在，流程结束。
  • 将 nums[i] 加到你的累计总和中。
  • 将索引 i 标记为已使用，并将 step 增加 1。

返回通过选择索引来获得的 最大总和。

 

示例 1：

输入：nums = [1,10,4,2,1,6], threshold = [5,1,5,5,2,2]

输出：17

解释：

  • 在 step = 1，由于 threshold[1] <= step，选择 i = 1。总和变为 10。标记下标 1。
  • 在 step = 2，由于 threshold[4] <= step，选择 i = 4。总和变为 11。标记下标 4。
  • 在 step = 3，由于 threshold[5] <= step，选择 i = 5。总和变为 17。标记下标 5。
  • 在 step = 4，我们不能选择下标 0，2 或 3 因为它们的阈值 > 4，所以我们结束流程。

示例 2：

输入：nums = [4,1,5,2,3], threshold = [3,3,2,3,3]

输出：0

解释：

在 step = 1 时没有下标 i 使得 threshold[i] <= 1，所以流程立刻结束。因此，总和为 0。

示例 3：

输入：nums = [2,6,10,13], threshold = [2,1,1,1]

输出：31

解释：

  • 在 step = 1，由于 threshold[3] <= step，选择 i = 3。总和变为 13。标记下标 3。
  • 在 step = 2，由于 threshold[2] <= step，选择 i = 2。总和变为 23。标记下标 2。
  • 在 step = 3，由于 threshold[1] <= step，选择 i = 1。总和变为 29。标记下标 1。
  • 在 step = 4，由于 threshold[0] <= step，选择 i = 0。总和变为 31。标记下标 0。
  • 在 step = 4 后所有下标都已经被选择，所以流程结束。

 

提示：

  • n == nums.length == threshold.length
  • 1 <= n <= 105
  • 1 <= nums[i] <= 109
  • 1 <= threshold[i] <= n

解法

方法一：贪心 + 排序

我们注意到，在每一个步骤中，我们都希望选择一个满足条件的数中最大的数加入总和中。因此，我们可以使用贪心的方法来解决这个问题。

我们首先将长度为 \(n\) 的下标数组 \(\textit{idx}\) 按照对应的阈值从小到大进行排序。然后，我们使用有序集合或优先队列（最大堆）来维护当前满足条件的数。在每一个步骤中，我们将所有阈值小于等于当前步骤数的数加入有序集合或优先队列中，然后选择其中最大的数加入总和中。如果此时有序集合或优先队列为空，说明没有满足条件的数，我们就结束过程。

时间复杂度 \(O(n \times \log n)\)，空间复杂度 \(O(n)\)。其中 \(n\) 为数组 \(\textit{nums}\) 的长度。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
class Solution:
    def maxSum(self, nums: List[int], threshold: List[int]) -> int:
        n = len(nums)
        idx = sorted(range(n), key=lambda i: threshold[i])
        sl = SortedList()
        step = 1
        ans = i = 0
        while True:
            while i < n and threshold[idx[i]] <= step:
                sl.add(nums[idx[i]])
                i += 1
            if not sl:
                break
            ans += sl.pop()
            step += 1
        return ans

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
class Solution {
    public long maxSum(int[] nums, int[] threshold) {
        int n = nums.length;
        Integer[] idx = new Integer[n];
        Arrays.setAll(idx, i -> i);
        Arrays.sort(idx, Comparator.comparingInt(i -> threshold[i]));
        TreeMap<Integer, Integer> tm = new TreeMap<>();
        long ans = 0;
        for (int i = 0, step = 1;; ++step) {
            while (i < n && threshold[idx[i]] <= step) {
                tm.merge(nums[idx[i]], 1, Integer::sum);
                ++i;
            }
            if (tm.isEmpty()) {
                break;
            }
            int x = tm.lastKey();
            ans += x;
            if (tm.merge(x, -1, Integer::sum) == 0) {
                tm.remove(x);
            }
        }
        return ans;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
class Solution {
public:
    long long maxSum(vector<int>& nums, vector<int>& threshold) {
        int n = nums.size();
        vector<int> idx(n);
        iota(idx.begin(), idx.end(), 0);
        sort(idx.begin(), idx.end(), [&](int a, int b) { return threshold[a] < threshold[b]; });

        multiset<int> ms;
        long long ans = 0;
        int i = 0;

        for (int step = 1;; ++step) {
            while (i < n && threshold[idx[i]] <= step) {
                ms.insert(nums[idx[i]]);
                ++i;
            }
            if (ms.empty()) {
                break;
            }

            auto it = prev(ms.end());
            ans += *it;
            ms.erase(it);
        }
        return ans;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
func maxSum(nums []int, threshold []int) int64 {
    n := len(nums)
    idx := make([]int, n)
    for i := 0; i < n; i++ {
        idx[i] = i
    }
    sort.Slice(idx, func(a, b int) bool {
        return threshold[idx[a]] < threshold[idx[b]]
    })

    tree := redblacktree.NewWithIntComparator()
    var ans int64
    i := 0

    for step := 1; ; step++ {
        for i < n && threshold[idx[i]] <= step {
            val := nums[idx[i]]
            if cnt, found := tree.Get(val); found {
                tree.Put(val, cnt.(int)+1)
            } else {
                tree.Put(val, 1)
            }
            i++
        }
        if tree.Empty() {
            break
        }

        node := tree.Right()
        key := node.Key.(int)
        cnt := node.Value.(int)

        ans += int64(key)
        if cnt == 1 {
            tree.Remove(key)
        } else {
            tree.Put(key, cnt-1)
        }
    }

    return ans
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
function maxSum(nums: number[], threshold: number[]): number {
    const n = nums.length;
    const idx = Array.from({ length: n }, (_, i) => i).sort((a, b) => threshold[a] - threshold[b]);
    const pq = new MaxPriorityQueue<number>();
    let ans = 0;
    for (let i = 0, step = 1; ; ++step) {
        while (i < n && threshold[idx[i]] <= step) {
            pq.enqueue(nums[idx[i]]);
            ++i;
        }
        if (pq.isEmpty()) {
            break;
        }
        ans += pq.dequeue();
    }
    return ans;
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
use std::cmp::Reverse;
use std::collections::BinaryHeap;

impl Solution {
    pub fn max_sum(nums: Vec<i32>, threshold: Vec<i32>) -> i64 {
        let n = nums.len();
        let mut idx: Vec<usize> = (0..n).collect();
        idx.sort_by_key(|&i| threshold[i]);

        let mut pq = BinaryHeap::new();
        let mut ans: i64 = 0;
        let mut i = 0;
        let mut step = 1;

        loop {
            while i < n && threshold[idx[i]] <= step {
                pq.push(nums[idx[i]]);
                i += 1;
            }
            match pq.pop() {
                Some(x) => {
                    ans += x as i64;
                    step += 1;
                }
                None => break,
            }
        }

        ans
    }
}

评论