1090. Largest Values From Labels

We have a set of items: the i-th item has value values[i] and label labels[i].

Then, we choose a subset S of these items, such that:

  • |S| <= num_wanted
  • For every label L, the number of items in S with label L is <= use_limit.

Return the largest possible sum of the subset S.

 

Example 1:

Input: values = [5,4,3,2,1], labels = [1,1,2,2,3], num_wanted = 3, use_limit = 1
Output: 9
Explanation: The subset chosen is the first, third, and fifth item.

Example 2:

Input: values = [5,4,3,2,1], labels = [1,3,3,3,2], num_wanted = 3, use_limit = 2
Output: 12
Explanation: The subset chosen is the first, second, and third item.

Example 3:

Input: values = [9,8,8,7,6], labels = [0,0,0,1,1], num_wanted = 3, use_limit = 1
Output: 16
Explanation: The subset chosen is the first and fourth item.

Example 4:

Input: values = [9,8,8,7,6], labels = [0,0,0,1,1], num_wanted = 3, use_limit = 2
Output: 24
Explanation: The subset chosen is the first, second, and fourth item.

 

Note:

  1. 1 <= values.length == labels.length <= 20000
  2. 0 <= values[i], labels[i] <= 20000
  3. 1 <= num_wanted, use_limit <= values.length

Rust Solution

struct Solution;
use std::collections::HashMap;
type Pair = (i32, i32);

impl Solution {
    fn largest_vals_from_labels(
        values: Vec<i32>,
        labels: Vec<i32>,
        num_wanted: i32,
        use_limit: i32,
    ) -> i32 {
        let n = values.len();
        let use_limit = use_limit as usize;
        let mut num_wanted = num_wanted as usize;
        let mut pairs: Vec<Pair> = values.into_iter().zip(labels.into_iter()).collect();
        pairs.sort_unstable();
        let mut hm: HashMap<i32, usize> = HashMap::new();
        let mut res = 0;
        for i in (0..n).rev() {
            let count = hm.entry(pairs[i].1).or_default();
            if *count < use_limit {
                *count += 1;
                res += pairs[i].0;
                num_wanted -= 1;
            }
            if num_wanted == 0 {
                break;
            }
        }
        res
    }
}

#[test]
fn test() {
    let values = vec![5, 4, 3, 2, 1];
    let labels = vec![1, 1, 2, 2, 3];
    let num_wanted = 3;
    let use_limit = 1;
    let res = 9;
    assert_eq!(
        Solution::largest_vals_from_labels(values, labels, num_wanted, use_limit),
        res
    );
    let values = vec![5, 4, 3, 2, 1];
    let labels = vec![1, 3, 3, 3, 2];
    let num_wanted = 3;
    let use_limit = 2;
    let res = 12;
    assert_eq!(
        Solution::largest_vals_from_labels(values, labels, num_wanted, use_limit),
        res
    );
    let values = vec![9, 8, 8, 7, 6];
    let labels = vec![0, 0, 0, 1, 1];
    let num_wanted = 3;
    let use_limit = 1;
    let res = 16;
    assert_eq!(
        Solution::largest_vals_from_labels(values, labels, num_wanted, use_limit),
        res
    );
    let values = vec![9, 8, 8, 7, 6];
    let labels = vec![0, 0, 0, 1, 1];
    let num_wanted = 3;
    let use_limit = 2;
    let res = 24;
    assert_eq!(
        Solution::largest_vals_from_labels(values, labels, num_wanted, use_limit),
        res
    );
}

Having problems with this solution? Click here to submit an issue on github.