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.PriorityQueue
@Suppress("NAME_SHADOWING")
class Solution {
private class Node(var `val`: Int, var label: Int)
fun largestValsFromLabels(values: IntArray, labels: IntArray, numWanted: Int, useLimit: Int): Int {
var numWanted = numWanted
val maxHeap =
PriorityQueue { a: Node, b: Node -> if (b.`val` != a.`val`) b.`val` - a.`val` else a.label - b.label }
val n = values.size
for (i in 0 until n) {
maxHeap.offer(Node(values[i], labels[i]))
}
var ans = 0
val labelAddedCount: HashMap<Int, Int> = HashMap()
while (maxHeap.isNotEmpty() && numWanted > 0) {
val cur = maxHeap.poll()
if (labelAddedCount.containsKey(cur.label) &&
labelAddedCount[cur.label]!! >= useLimit
) {
continue
}
if (cur.`val` > 0) {
ans += cur.`val`
labelAddedCount[cur.label] = labelAddedCount.getOrDefault(cur.label, 0) + 1
numWanted--
}
}
return ans
}
}