Medium
There is a set of n
items. You are given two integer arrays values
and labels
where the value and the label of the ith
element are values[i]
and labels[i]
respectively. You are also given two integers numWanted
and useLimit
.
Choose a subset s
of the n
elements such that:
s
is less than or equal to numWanted
.useLimit
items with the same label in s
.The score of a subset is the sum of the values in the subset.
Return the maximum score of a subset s
.
Example 1:
Input: values = [5,4,3,2,1], labels = [1,1,2,2,3], numWanted = 3, useLimit = 1
Output: 9
Explanation: The subset chosen is the first, third, and fifth items.
Example 2:
Input: values = [5,4,3,2,1], labels = [1,3,3,3,2], numWanted = 3, useLimit = 2
Output: 12
Explanation: The subset chosen is the first, second, and third items.
Example 3:
Input: values = [9,8,8,7,6], labels = [0,0,0,1,1], numWanted = 3, useLimit = 1
Output: 16
Explanation: The subset chosen is the first and fourth items.
Constraints:
n == values.length == labels.length
1 <= n <= 2 * 104
0 <= values[i], labels[i] <= 2 * 104
1 <= numWanted, useLimit <= n
import java.util.HashMap;
import java.util.PriorityQueue;
public class Solution {
private static class Node {
int val;
int label;
Node(int val, int label) {
this.val = val;
this.label = label;
}
}
public int largestValsFromLabels(int[] values, int[] labels, int numWanted, int useLimit) {
PriorityQueue<Node> maxHeap =
new PriorityQueue<>((a, b) -> b.val != a.val ? b.val - a.val : a.label - b.label);
int n = values.length;
for (int i = 0; i < n; i++) {
maxHeap.offer(new Node(values[i], labels[i]));
}
int ans = 0;
HashMap<Integer, Integer> labelAddedCount = new HashMap<>();
while (!maxHeap.isEmpty() && numWanted > 0) {
Node cur = maxHeap.poll();
if (labelAddedCount.containsKey(cur.label)
&& labelAddedCount.get(cur.label) >= useLimit) {
continue;
}
if (cur.val > 0) {
ans += cur.val;
labelAddedCount.put(cur.label, labelAddedCount.getOrDefault(cur.label, 0) + 1);
numWanted--;
}
}
return ans;
}
}