Problem
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
解题思路
根据题意可以知道,我们需要在每个对应的 l a b e l label label中选择不超过 u s e _ l i m i t use\_limit use_limit个数字,总的数字数目不超过 n u m _ w a n t e d num\_wanted num_wanted,求这些数字的可以组成的最大的和。那么我们只需要对每个 l a b e l label label对应的数据进行排序,挑选出最大的 u s e _ l i m i t use\_limit use_limit 个数字,将这些数字放进同一个数组中,进行排序,在挑选出最大的 n u m _ w a n t e d num\_wanted num_wanted个数字进行求和即可。由于这里需要进行排序挑选,因此也可以使用heap进行求解,原理是一致的。
代码如下:
class Solution:
def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
record = collections.defaultdict(list)
for v, l in zip(values, labels):
record[l].append(v)
seq = []
for vs in record.values():
seq += heapq.nlargest(min(len(vs), use_limit), vs)
return sum(heapq.nlargest(min(len(seq), num_wanted), seq))
或者:
class Solution:
def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
record = collections.defaultdict(list)
for v, l in zip(values, labels):
record[l].append(v)
seq = []
for vs in record.values():
seq += sorted(vs, reverse=True)[:min(len(vs), use_limit)]
seq.sort(reverse=True)
return sum(seq[:min(len(seq), num_wanted)])