跳转至

3763. Maximum Total Sum with Threshold Constraints 🔒

题目描述

You are given two integer arrays nums and threshold, both of length n.

Starting at step = 1, you perform the following repeatedly:

  • Choose an unused index i such that threshold[i] <= step.
    • If no such index exists, the process ends.
  • Add nums[i] to your running total.
  • Mark index i as used and increment step by 1.

Return the maximum total sum you can obtain by choosing indices optimally.

 

Example 1:

Input: nums = [1,10,4,2,1,6], threshold = [5,1,5,5,2,2]

Output: 17

Explanation:

  • At step = 1, choose i = 1 since threshold[1] <= step. The total sum becomes 10. Mark index 1.
  • At step = 2, choose i = 4 since threshold[4] <= step. The total sum becomes 11. Mark index 4.
  • At step = 3, choose i = 5 since threshold[5] <= step. The total sum becomes 17. Mark index 5.
  • At step = 4, we cannot choose indices 0, 2, or 3 because their thresholds are > 4, so we end the process.

Example 2:

Input: nums = [4,1,5,2,3], threshold = [3,3,2,3,3]

Output: 0

Explanation:

At step = 1 there is no index i with threshold[i] <= 1, so the process ends immediately. Thus, the total sum is 0.

Example 3:

Input: nums = [2,6,10,13], threshold = [2,1,1,1]

Output: 31

Explanation:

  • At step = 1, choose i = 3 since threshold[3] <= step. The total sum becomes 13. Mark index 3.
  • At step = 2, choose i = 2 since threshold[2] <= step. The total sum becomes 23. Mark index 2.
  • At step = 3, choose i = 1 since threshold[1] <= step. The total sum becomes 29. Mark index 1.
  • At step = 4, choose i = 0 since threshold[0] <= step. The total sum becomes 31. Mark index 0.
  • After step = 4 all indices have been chosen, so the process ends.

 

Constraints:

  • 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
    }
}

评论