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:

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

Solution:

class Solution {
    public int largestValsFromLabels(int[] values, int[] labels, int num_wanted, int use_limit) {
        int n = values.length;
        int[][] nodes = new int[n][2];
        for (int i = 0; i < n; i ++) {
            nodes[i] = new int[]{values[i], labels[i]};
        }
        Arrays.sort(nodes, (a, b) -> b[0] - a[0]);
        int count = 0, sum = 0;
        int[] map = new int[20001];
        for (int i = 0; i < n; i ++) {
            int[] node = nodes[i];
            if (map[node[1]] < use_limit) {
                sum += node[0];
                map[node[1]] ++;
                count ++;
            }
            if (count == num_wanted) break;
        }
        return sum;
    }
}