LeetCode Q 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 <= values.length == labels.length <= 20000
0 <= values[i], labels[i] <= 20000
1 <= num_wanted, use_limit <= values.length

Solution

Use an array pair to store the value - label pairs;
Sort pair according to value;
Use an array count to store the number of labels we have used

Code:

public int largestValsFromLabels(int[] values, int[] labels, int num_wanted, int use_limit) {
  
  if (values.length == 0) return 0;

  int len = values.length;
  int[][] pairs = new int[len][2];

  for (int i = 0; i < len; i++) {
    pairs[i][0] = values[i]; pairs[i][1] = labels[i];
  }

  Arrays.sort(pairs, (a, b) -> (b[0] - a[0]));

  int[] count = new int[20001];
  int totalNumber = 0, totalSum = 0;

  for (int i = 0; i < len; i++) {
    if (totalNumber == num_wanted) break;

    if (count[pairs[i][1]] == use_limit) continue;

    totalSum += pairs[i][0];
    count[pairs[i][1]]++;
    totalNumber++;
  }

  return totalSum;
}

Reprint policy