Largest Values From Labels in Python

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.

Constraints:

 1 <= values.length == labels.length <= 20000

 0 <= values[i], labels[i] <= 20000

 1 <= num_wanted, use_limit <= values.length 

Solution:

class Solution(object):
    def largestValsFromLabels(values, labels, num_wanted, use_limit):
        sorted_val = sorted([(i, j) for i, j in zip(values, labels)], key = lambda x : x[0]*-1)
        label_count = {label: 0 for label in set(labels)}
        result = 0
        
        for s_v in sorted_val:
            if num_wanted:
                if label_count[s_v[1]] < use_limit:
                    result += s_v[0]
                    label_count[s_v[1]] +=1
                    num_wanted -= 1
            else:
                break
        return result


# answer 1
values, labels, num_wanted, use_limit  = [5,4,3,2,1], [1,1,2,2,3], 3, 1
print(Solution.largestValsFromLabels(values, labels, num_wanted, use_limit))

# answer 2
values, labels, num_wanted, use_limit  = [5,4,3,2,1], [1,3,3,3,2], 3, 2
print(Solution.largestValsFromLabels(values, labels, num_wanted, use_limit))

# answer 3
values, labels, num_wanted, use_limit  = [9,8,8,7,6], [0,0,0,1,1], 3, 1
print(Solution.largestValsFromLabels(values, labels, num_wanted, use_limit))

# answer 4
values, labels, num_wanted, use_limit  = [9,8,8,7,6], [0,0,0,1,1], 3, 2
print(Solution.largestValsFromLabels(values, labels, num_wanted, use_limit))