3767 - Maximize Points After Choosing K Tasks

Description (leetcode)

You are given two integer arrays, technique1 and technique2, each of length n, where n represents the number of tasks to complete.

• If the i^th task is completed using technique 1, you earn technique1[i] points.
• If it is completed using technique 2, you earn technique2[i] points.

You are also given an integer k, representing the minimum number of tasks that must be completed using technique 1.

You must complete at least k tasks using technique 1 (they do not need to be the first k tasks).

The remaining tasks may be completed using either technique.

Return an integer denoting the maximum total points you can earn.

Example 1:

Input: technique1 = [5,2,10], technique2 = [10,3,8], k = 2

Output: 22

Explanation:

We must complete at least k = 2 tasks using technique1.

Choosing technique1[1] and technique1[2] (completed using technique 1), and technique2[0] (completed using technique 2), yields the maximum points: 2 + 10 + 10 = 22.

Example 2:

Input: technique1 = [10,20,30], technique2 = [5,15,25], k = 2

Output: 60

Explanation:

We must complete at least k = 2 tasks using technique1.

Choosing all tasks using technique 1 yields the maximum points: 10 + 20 + 30 = 60.

Example 3:

Input: technique1 = [1,2,3], technique2 = [4,5,6], k = 0

Output: 15

Explanation:

Since k = 0, we are not required to choose any task using technique1.

Choosing all tasks using technique 2 yields the maximum points: 4 + 5 + 6 = 15.

Constraints:

• 1 <= n == technique1.length == technique2.length <= 10⁵
• 1 <= technique1[i], technique2​​​​​​​[i] <= 10^{​​​​​​​5}
• 0 <= k <= n

submission

impl Solution {
    pub fn max_points(
        technique1: Vec<i32>,
        technique2: Vec<i32>,
        k: i32,
    ) -> i64 {
        // quick selection is much faster than our previous BinaryHeap solution.
        if technique1.len() == k as usize {
            // index guard
            return technique1.into_iter().map(|t| t as i64).sum::<i64>();
        }
        let mut techniques: Vec<_> = technique1
            .into_iter()
            .zip(technique2)
            // we want the fisrt k task to maximize tech1 points
            .map(|(t1, t2)| ((t2 - t1) as i64, t2 as i64))
            .collect();
        let (lesser, median, greater) =
            techniques.select_nth_unstable(k as usize);
        // k tech1 tasks
        lesser.iter_mut().map(|&mut (x1, x2)| -x1 + x2).sum::<i64>()
            // pick whichever is larger for the rest
            + std::iter::once(median)
                .chain(greater.iter_mut())
                .map(|&mut (x1, x2)| x2.max(-x1 + x2))
                .sum::<i64>()
    }
}