Skip to content

3763. Maximum Total Sum with Threshold Constraints πŸ”’

Description

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

Solutions

Solution 1: Greedy + Sorting

We observe that at each step, we want to select the largest number among those that satisfy the condition to add to the total sum. Therefore, we can use a greedy approach to solve this problem.

We first sort an index array \(\textit{idx}\) of length \(n\) in ascending order by their corresponding thresholds. Then, we use a sorted set or priority queue (max heap) to maintain the numbers that currently satisfy the condition. At each step, we add all numbers whose thresholds are less than or equal to the current step number into the sorted set or priority queue, and then select the largest number among them to add to the total sum. If the sorted set or priority queue is empty at this point, it means there are no numbers that satisfy the condition, and we end the process.

The time complexity is \(O(n \times \log n)\), and the space complexity is \(O(n)\), where \(n\) is the length of the array \(\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
    }
}

Comments