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

Difficulty: interview

Starter code:

class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        

Solutions:
class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        
        
        
        # my solution ... 128 ms ... 99 % ... 17.9 MB ... 85 %
        #  time: O(nlogn)
        # space: O(n)
        
        l2v = collections.defaultdict(list)
        for v,l in zip(values, labels):
            l2v[l].append(v)
        pool = []
        for l in l2v:
            pool += sorted(l2v[l])[-use_limit:]
        return sum(sorted(pool)[-num_wanted:])
        
        


class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
            val_label = list(sorted(zip(values, labels), key = lambda x: x[0], reverse = True))

            cnt_label  = collections.Counter()
            chosen = 0
            res = 0

            
            for val, lab in val_label:
                if chosen < num_wanted:
                    if cnt_label[lab] < use_limit:
                        res += val
                        chosen += 1
                        cnt_label[lab] += 1
            return res



class Item:
    def __init__(self,label,value):
        self.label=label
        self.value=value

class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        labelvalues=[]
        i=0
        while i <len(values):
            labelvalues.append(Item(labels[i],values[i]))
            i+=1
        lv=[]
        
        for item in labelvalues:
            heapq.heappush(lv,(-item.value,item.label))
        countlabels=collections.defaultdict(int)
        maximum_sum=0
        while num_wanted>0 and lv:
            max_value=heapq.heappop(lv)
            curr_max=max_value[0]
            curr_label_value=max_value[1]
            if countlabels.get(curr_label_value,0)<use_limit:
                countlabels[curr_label_value]+=1
                maximum_sum-=curr_max
                num_wanted-=1
        return maximum_sum

class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        res = 0
        heap = [[-v, i] for v, i in zip(values, labels)]
        heapq.heapify(heap)
        counter = Counter()
        
        while heap and num_wanted:
            v, i = heapq.heappop(heap)
            if counter[i] == use_limit:
                continue
            
            counter[i] += 1
            res += -v
            num_wanted -= 1
            
        return res

class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
#         # maintain max heap, get largest value for each label
#         # this brute force solution is O(nlogn + num_wanted)
#         values_by_label = defaultdict(list)
#         for i, val in enumerate(values):
#                 heapq.heappush(values_by_label[labels[i]], -1 * val)
        
#         total = 0
#         label_use_limit = defaultdict(int)
#         for label in values_by_label.keys():
#             label_use_limit[label] = use_limit

            
#         for _ in range(num_wanted - 1, -1, -1):
#             max_num = -1
#             for k, v in values_by_label.items():
#                 if label_use_limit[k] > 0 and len(v) > 0 and -1 * v[0] > max_num:
#                     max_num = -1 * v[0]
#                     best_label = k  
#             if max_num != -1:
#                 total += max_num
#                 heapq.heappop(values_by_label[best_label])
#                 label_use_limit[best_label] -= 1
        
#         return total

        ordered_values = [] # maxheap
        label_use_limit = defaultdict(int)
        for i, val in enumerate(values):
            heapq.heappush(ordered_values, (-1 * val, labels[i]))
            if labels[i] not in label_use_limit:
                label_use_limit[labels[i]] = use_limit
        total = 0
        while len(ordered_values) > 0 and num_wanted > 0:
            top_elem = heapq.heappop(ordered_values)
            if label_use_limit[top_elem[1]] > 0:
                total += (-1 * top_elem[0])
                label_use_limit[top_elem[1]] -= 1
                num_wanted -= 1
        return total


class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        label2values = defaultdict(list)
        for value, label in zip(values, labels):
            label2values[label].append(value)
            
        all_top_nums = []
        for values in list(label2values.values()):
            if use_limit >= len(values):
                top_nums = values
            else:
                top_nums = self.quick_select(use_limit, values)
            all_top_nums.extend(top_nums)
            
        if num_wanted >= len(all_top_nums):
            return sum(all_top_nums)
        return sum(self.quick_select(num_wanted, all_top_nums))
    
    def quick_select(self, use_limit, nums):
        low = 0
        high = len(nums) - 1
        use_limit -= 1
        while low <= high:
            pivot = self.partition(nums, low, high)
            if pivot == use_limit:
                return nums[:pivot + 1]
            elif pivot < use_limit:
                low = pivot + 1
            else:
                high = pivot - 1
        return nums
        
    def partition(self, nums, low, high):
        pivot = random.randint(low, high)
        self._swap(nums, pivot, high)
        pivot = low
        for i in range(low, high):
            if nums[i] > nums[high]:
                self._swap(nums, pivot, i)
                pivot += 1
        self._swap(nums, pivot, high)
        return pivot
    
    def _swap(self, nums, i, j):
        nums[i], nums[j] = nums[j], nums[i]
        
# Follow-up 1
# Do beter than O(nlogn).

# Follow-up 2
# Can you do it in O(n) time?

# What about using QuickSelect?

# Steps:

# Separate all elements by label => HashMap<Label, Array> m;
# Use QuickSelect on each label's elements to find the largest m elements.
# Merge the largest m elements from each label into an array
# Use QuickSelect again on the resulting array to get the largest k elements.
# *each step can be done in O(N).


import heapq 

class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        
        maxHeap = []
        
        for value, label in zip(values, labels):
            print((value, label))
            heappush(maxHeap, (-value, label))
            
        labelFrequencies = {} 
        maxSum = 0
        numChosen = 0
        
        while maxHeap and numChosen < num_wanted: 
            value, label = heappop(maxHeap)
            value = -value
            
            #when you can't use that value 
            if label in labelFrequencies and labelFrequencies[label] >= use_limit:
                continue 
            
            if label not in labelFrequencies:
                labelFrequencies[label] = 0 
                
            labelFrequencies[label] += 1 
            numChosen += 1 
            maxSum += value 
            
        return maxSum 
                
            


class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        # Greedy search
        # 1. sorted each item based on used frequency
        # 2. loop through and items until bag full
        
        items = sorted(list(zip(values, labels)), key=lambda x: -x[0])
        
        res = 0
        freq = defaultdict(int)
        n = 0
        for i, (val, label) in enumerate(items):
            if freq[label] < use_limit:
                freq[label] += 1
                res += val
                n += 1
            
            if n == num_wanted:
                break
        
        return res
            
            


from collections import defaultdict

class Solution:
    def quickselect(self, nums, K):
        def swap(nums, i, j):
            nums[i], nums[j] = nums[j], nums[i]

        def partition(nums, left, right, pivotIdx):
            if right-left <= 1:
                return left
            
            pivot = nums[pivotIdx]
            swap(nums, pivotIdx, right-1)
            
            read_idx = left
            write_idx = left
            while read_idx < right-1:
                if nums[read_idx] > pivot:
                    swap(nums, read_idx, write_idx)
                    write_idx += 1
                read_idx += 1
            
            swap(nums, write_idx, right-1)
            return write_idx
        
        left = 0
        right = len(nums)
        pivotIdx = -1
        while pivotIdx != K:
            pivotIdx = left + (right-left)//2
            pivotIdx = partition(nums, left, right, pivotIdx)
            
            if pivotIdx < K:
                left = pivotIdx+1
            elif pivotIdx > K:
                right = pivotIdx
                

    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        valuesByLabel = defaultdict(list)
        for value, label in zip(values, labels):
            valuesByLabel[label].append(value)
            
        valuesWithinLimit = []
        for i_label, i_values in valuesByLabel.items():
            self.quickselect(i_values, use_limit)
            valuesWithinLimit += i_values[:use_limit]
            
        self.quickselect(valuesWithinLimit, num_wanted)
        return sum(valuesWithinLimit[:num_wanted])

class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        dic = {}
        for i,j in zip(values,labels):
            if j not in dic:
                dic[j] = [i]
            else:
                dic[j].append(i)
        
        res_arr = []
        
        for i in dic:
            res_arr.extend(sorted(dic[i],reverse=True)[0:use_limit])
        
        res_arr.sort()
        count = 0
        res = 0
        while count < num_wanted and len(res_arr) > 0:
            tmp  = res_arr.pop()
            res += tmp
            count += 1
        return res
        


class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        res = 0
        arr = sorted(list(zip(values, labels)), key=lambda x: x[0], reverse=1)
        cnt_label = {}
        cnt_tot = 0
        for val, lab in arr:
            if cnt_tot < num_wanted and cnt_label.get(lab, 0) < use_limit:
                cnt_tot += 1
                cnt_label[lab] = cnt_label.get(lab, 0) + 1
                res += val
        return res


class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        d =defaultdict(list)
        n = len(labels)
        for i in range(n):
            d[labels[i]].append(values[i])
        
        arr = []
        for k,v in list(d.items()):
            d[k].sort(reverse=True)
            d[k] = d[k][:use_limit]
            arr.extend(d[k])
        
        arr.sort(reverse=True)
        return sum(arr[:num_wanted])
        
        
            


class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        heap = []
        for v, l in zip(values, labels):
            heapq.heappush(heap, (-v, l))
        
        count = collections.Counter()
        res = []
        while heap and len(res) < num_wanted:
            v, l = heapq.heappop(heap)
            if count[l] < use_limit:
                res.append(v)
                count[l] += 1
        return -sum(res)
        


from collections import defaultdict

class Solution:

    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        
        if len(values) == 0:
            return 0
        if len(values) != len(labels):
            raise ValueError
        
        sorted_value_labels = sorted(zip(values, labels), key=lambda x: -x[0])
        label_counter = defaultdict(lambda: 0)
        
        i = 0
        total = 0
        S = 0

        while S != num_wanted and i < len(values):
            val, label = sorted_value_labels[i]
            if label_counter[label] < use_limit:
                total += val
                label_counter[label] += 1
                S += 1
            i += 1
        
        return total
                
            
            
        


class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        sum = 0
        ct = {}
        import numpy as np
        sorted_idx = np.argsort(-np.array(values))
        select_ct = 0
        for v in sorted_idx:
            print(v)
            if select_ct == num_wanted:
                break
            label = labels[v]
            if label not in list(ct.keys()):
                ct[label] = 1
                sum += values[v]
                select_ct += 1
            elif ct[label] < use_limit:
                ct[label] += 1
                sum += values[v]
                select_ct += 1
            else:
                continue
            print(sum)
        return sum
        


class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        newList = [(v, l) for v, l in zip(values, labels)]
        newList.sort(reverse = True)
        
        ans = 0
        i = 0
        labelCounter = defaultdict(int)
        for v, l in newList:
            if labelCounter[l] < use_limit:
                # print(v, l)
                ans += v
                i += 1
                labelCounter[l] += 1
                if i == num_wanted:
                    break
        return ans

from collections import defaultdict

class Solution:
    def largestValsFromLabels(self, values, labels, num_wanted, use_limit):
        if not values or not labels:
            return 0
        index = [i for i in range(len(values))]
        index.sort(key=lambda i: values[i], reverse=True)
        res = 0
        lab_to_count = defaultdict(int)
        for i in index:
            if num_wanted == 0:
                return res
            l = labels[i]
            v = values[i]
            if l not in lab_to_count or lab_to_count[l] < use_limit:
                res += v
                lab_to_count[l] += 1
                num_wanted -= 1    
        return res


class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        all_value_labels = [(value, label) for value, label in zip(values, labels)]
        all_value_labels.sort(reverse=True)
        used = Counter()
        ans = []
        for value, label in all_value_labels:
            if used[label] + 1 <= use_limit:
                used[label] += 1
                ans.append(value)
                if len(ans) == num_wanted:
                    break
        return sum(ans)


class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        heap = []
        for i in range(len(values)):
            heappush(heap, (-values[i], labels[i]))
        selected, label_counts = [], {}
        while len(selected) < num_wanted and heap:
            value, label = heappop(heap)
            if (cur_num := label_counts.get(label, 0)) < use_limit:
                label_counts[label] = cur_num + 1
                selected.append(value)
        print(selected)
        return -sum(selected)

class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        counts = collections.defaultdict(int)
        vl = sorted(zip(values,labels))
        ans = 0
        while num_wanted and vl:
            val,lab = vl.pop()
            if counts[lab] < use_limit:
                ans += val
                counts[lab] += 1
                num_wanted -= 1
        return ans

from collections import defaultdict

class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        
        counts = defaultdict(int)
        s = 0
        total_count = 0
        
        inds = sorted(list(range(len(values))), key=lambda x: values[x], reverse=True)
        
        sorted_values = [values[idx] for idx in inds]
        sorted_labels = [labels[idx] for idx in inds]
        
        for i, val in enumerate(sorted_values):
            label = sorted_labels[i]
            if total_count < num_wanted and counts[label] < use_limit:
                counts[label] += 1
                total_count += 1
                s += val
                
            if total_count == num_wanted:
                break
        return s


class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        vls = sorted(zip(values, labels), reverse=True)
        l2v = defaultdict(list)
        for v,l in vls:
            l2v[l].append(v)
        chosen = []
        l_use = defaultdict(int)
        for i, (v, l) in enumerate(vls):
            if l_use[l] < use_limit:
                chosen.append((v,l))
                l_use[l] += 1
                if len(chosen) >= num_wanted:
                    break
        # print(chosen)
        return sum(v for v,l in chosen)

from collections import Counter
import heapq

class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        h = list(zip([-v for v in values], labels))

        heapq.heapify(h)
        
        counter = Counter()

        ans = 0
        while num_wanted and h:
            heap_item = heapq.heappop(h)
            val, label = heap_item
            if counter[label] == use_limit:
                continue
                
            counter[label] += 1
            ans += -val
            num_wanted -= 1
            
        return ans


class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        items = []
        for i in range(len(values)):
            items.append((values[i], labels[i]))
    
        items.sort(key = lambda item: item[0], reverse=True)
        dic = Counter(labels)

        
        for item in dic:
            if dic[item] > use_limit:
                dic[item] = use_limit
                
        largest = 0
        print(dic)
        print(items)
        for num in items:
            if num_wanted == 0:
                break
            current_label = num[1]
            value = num[0]
            print(current_label)
            
            if dic[current_label] > 0:
                num_wanted -= 1
                dic[current_label] -= 1
                largest += value
        return largest


from queue import PriorityQueue

class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        pq = PriorityQueue()
        for i in range(len(values)):
            pq.put((-values[i],labels[i]))
        
        subset_sum = 0
        ecount = 0
        label_usage_count_map = {}
        
        while (pq.qsize()>0):
            popped_element = pq.get()
            v = -popped_element[0]
            l = popped_element[1]
            
            if (v<0):
                break
            
            if (ecount>=num_wanted):
                break
            
            lusage = label_usage_count_map.get(l,0)
            if (lusage>=use_limit):
                continue
            
            ecount += 1
            subset_sum += v
            label_usage_count_map[l] = label_usage_count_map.get(l,0) + 1
        
        return subset_sum


class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        empty = {i : use_limit for i in set(labels)}
        ans = 0
        pairs = [(values[i], labels[i]) for i in range(len(values))]
        pairs.sort(key = lambda x : (-x[0]))
        for index, pair in enumerate(pairs):
            if empty[pair[1]] > 0 and num_wanted > 0:
                ans += pair[0]
                empty[pair[1]] -= 1
                num_wanted -= 1
        return ans

class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        val_label_tuple = [(-1*val, lab) for val, lab in zip(values, labels)]
        import heapq
        heapq.heapify(val_label_tuple) #maximum val is at the top
        counts = collections.defaultdict(int)
        
        ans = 0
        while val_label_tuple:
            val, lab = heapq.heappop(val_label_tuple)
            if counts.get(lab, 0) < use_limit:
                ans += abs(val)
                num_wanted -= 1
                if num_wanted == 0:
                    break
                counts[lab] += 1
        
        return ans
                
        
        
        '''
        val_label_tuple = [(val, lab) for val, lab in zip(values, labels)]
        val_label_tuple = sorted(val_label_tuple, key=lambda x: x[0], reverse=True)
        
        count_items = defaultdict(int)
        ret_value = 0
        
        for val, label in val_label_tuple:
            if count_items.get(label, 0) < use_limit:
                ret_value += val
                count_items[label] += 1
                num_wanted -= 1
                if num_wanted == 0:
                    break
        
        return ret_value
        '''

from collections import defaultdict
from queue import PriorityQueue


class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        q = PriorityQueue()
        for i in range(len(values)):
            q.put((-values[i], values[i], labels[i]))
        d = defaultdict(lambda: 0)
        ans = []
        while len(ans) < num_wanted and not q.empty():
            _, n, l = q.get()
            if d[l] < use_limit:
                ans.append(n)
                d[l] += 1
        return sum(ans)

class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        values_w_labels = [(values[i], labels[i]) for i in range(len(values))]
        values_w_labels.sort(reverse=True)
        label_count = {}
        total = 0
        count = 0
        for value, label in values_w_labels:
            if label not in label_count:
                label_count[label] = 0
            if label_count[label] < use_limit:
                total += value
                label_count[label] += 1
                count += 1
            if count == num_wanted:
                break
        return total

from heapq import *
class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        maxHeap = []
        maxSum = 0
        
        for i in range(len(values)):
            heappush(maxHeap, (-values[i], labels[i]))
            
        usedLabels = defaultdict(int)
        
        while num_wanted > 0 and len(maxHeap) > 0:
            currVal, currLabel = heappop(maxHeap)
            
            if usedLabels[currLabel] < use_limit:
                usedLabels[currLabel] += 1
                maxSum -= currVal
                num_wanted -= 1
                
        return maxSum
                


class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        
        hashmap = collections.defaultdict(int)
        
        maxHeap = [(-v,labels[i]) for i,v in enumerate(values)]
        heapq.heapify(maxHeap)
        out = 0
        while num_wanted>0:
            if len(maxHeap)==0: 
                break 
            value, label = heapq.heappop(maxHeap)
            if hashmap[label]<use_limit:
                hashmap[label] = hashmap[label] +1 
                num_wanted -= 1
                out = out - value
        return out
        


class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        heap = []
        
        for i in range(len(values)):
            heappush(heap, (-values[i], labels[i]))
            
        if use_limit == 0 or num_wanted == 0:
            return 0
            
        used = {}
        ret = 0
        while len(heap) > 0:
            heap_item = heappop(heap)
            value = heap_item[0]
            label = heap_item[1]
            
            used[label] = used.get(label, 0) + 1
            if used[label] <= use_limit:
                ret += (-value)
                num_wanted -= 1
                
                if num_wanted == 0:
                    break
        
        return ret
            


class Solution:
    def largestValsFromLabels(self, values: List[int], labels: List[int], num_wanted: int, use_limit: int) -> int:
        from heapq import heapify, heappop
        from collections import Counter
        counter = Counter()
        pq, ans = [[-value, label] for value, label in zip(values, labels)], []
        heapify(pq)
        while pq:
            cur = heappop(pq)
            if counter[cur[1]] < use_limit:
                counter[cur[1]] += 1
                ans.append(-cur[0])
                if len(ans) == num_wanted:
                    break
        return sum(ans)