LeetCode in Kotlin

1090. Largest Values From Labels

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:

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:

Solution

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
    }
}