rustgym/leetcode/
_1090_largest_values_from_labels.rs

1struct Solution;
2use std::collections::HashMap;
3type Pair = (i32, i32);
4
5impl Solution {
6    fn largest_vals_from_labels(
7        values: Vec<i32>,
8        labels: Vec<i32>,
9        num_wanted: i32,
10        use_limit: i32,
11    ) -> i32 {
12        let n = values.len();
13        let use_limit = use_limit as usize;
14        let mut num_wanted = num_wanted as usize;
15        let mut pairs: Vec<Pair> = values.into_iter().zip(labels.into_iter()).collect();
16        pairs.sort_unstable();
17        let mut hm: HashMap<i32, usize> = HashMap::new();
18        let mut res = 0;
19        for i in (0..n).rev() {
20            let count = hm.entry(pairs[i].1).or_default();
21            if *count < use_limit {
22                *count += 1;
23                res += pairs[i].0;
24                num_wanted -= 1;
25            }
26            if num_wanted == 0 {
27                break;
28            }
29        }
30        res
31    }
32}
33
34#[test]
35fn test() {
36    let values = vec![5, 4, 3, 2, 1];
37    let labels = vec![1, 1, 2, 2, 3];
38    let num_wanted = 3;
39    let use_limit = 1;
40    let res = 9;
41    assert_eq!(
42        Solution::largest_vals_from_labels(values, labels, num_wanted, use_limit),
43        res
44    );
45    let values = vec![5, 4, 3, 2, 1];
46    let labels = vec![1, 3, 3, 3, 2];
47    let num_wanted = 3;
48    let use_limit = 2;
49    let res = 12;
50    assert_eq!(
51        Solution::largest_vals_from_labels(values, labels, num_wanted, use_limit),
52        res
53    );
54    let values = vec![9, 8, 8, 7, 6];
55    let labels = vec![0, 0, 0, 1, 1];
56    let num_wanted = 3;
57    let use_limit = 1;
58    let res = 16;
59    assert_eq!(
60        Solution::largest_vals_from_labels(values, labels, num_wanted, use_limit),
61        res
62    );
63    let values = vec![9, 8, 8, 7, 6];
64    let labels = vec![0, 0, 0, 1, 1];
65    let num_wanted = 3;
66    let use_limit = 2;
67    let res = 24;
68    assert_eq!(
69        Solution::largest_vals_from_labels(values, labels, num_wanted, use_limit),
70        res
71    );
72}
rustgym/leetcode/_1090_largest_values_from_labels.rs

rustgym/leetcode/
_1090_largest_values_from_labels.rs
