largest-values-from-labels.py

# Time:  O(nlogn)
# Space: O(n)

import collections


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
        """
        counts = collections.defaultdict(int)
        val_labs = zip(values,labels)
        val_labs.sort(reverse=True)
        result = 0
        for val, lab in val_labs:
            if counts[lab] >= use_limit:
                continue
            result += val
            counts[lab] += 1
            num_wanted -= 1
            if num_wanted == 0:
                break
        return result