LeetCode-Python-1090. 受标签影响的最大值

我们有一个项的集合，其中第 i 项的值为 values[i]，标签为 labels[i]。

我们从这些项中选出一个子集 S，这样一来：

• |S| <= num_wanted
• 对于任意的标签 L，子集 S 中标签为 L 的项的数目总满足 <= use_limit。

返回子集 S 的最大可能的 和。

 

示例 1：

输入：values = [5,4,3,2,1], labels = [1,1,2,2,3], num_wanted = 3, use_limit = 1
输出：9
解释：选出的子集是第一项，第三项和第五项。

示例 2：

输入：values = [5,4,3,2,1], labels = [1,3,3,3,2], num_wanted = 3, use_limit = 2
输出：12
解释：选出的子集是第一项，第二项和第三项。

示例 3：

输入：values = [9,8,8,7,6], labels = [0,0,0,1,1], num_wanted = 3, use_limit = 1
输出：16
解释：选出的子集是第一项和第四项。

示例 4：

输入：values = [9,8,8,7,6], labels = [0,0,0,1,1], num_wanted = 3, use_limit = 2
输出：24
解释：选出的子集是第一项，第二项和第四项。

 

提示：

1. 1 <= values.length == labels.length <= 20000
2. 0 <= values[i], labels[i] <= 20000
3. 1 <= num_wanted, use_limit <= values.length

思路：

题目有点难懂……实际意思是有很多个标签，每个标签里有很多个值，

题目给定限制use_limit，要求每个标签最多只能选use_limit个值，

问一共选num_wanted个值的最大总和是多少。

举例： values = [9,8,8,7,6], labels = [0,0,0,1,1], num_wanted = 3, use_limit = 2

对于标签0，它的值有[9, 8, 8]

对于标签1， 它的值有[7,6]

因为每个标签最多只能选两个，所以对于每个标签建立SIZE为2的最小堆，

可以得到标签0的堆为[9,8]， 标签1的堆为[7,6]，

然后题目就变成了在[9,8] + [7,6]里选最大的3个数字出来。

from heapq import *
class Solution(object):
    def largestValsFromLabels(self, values, labels, num_wanted, use_limit):
        """
        :type values: List[int]
        :type labels: List[int]
        :type num_wanted: int
        :type use_limit: int
        :rtype: int
        """
        record = dict()
        l = len(values)
        for i in range(l):
            if labels[i] not in record:
                record[labels[i]] = [values[i]]
                heapify(record[labels[i]])
            else:
                heappush(record[labels[i]], values[i])
                if len(record[labels[i]]) > use_limit:
                    heappop(record[labels[i]])
                    
        res = []
        for key, val in record.items():
            res += val
        res.sort()
        return sum(res[-num_wanted:])