Learning On the Way

347. Top K Frequent Elements

Given an integer array nums and an integer k, return the k most frequent elements. You may return the answer in any order.

Use bucket to search

def topKFrequent(self, nums: List[int], k: int) -> List[int]:
    from itertools import chain
    ct = collections.Counter(nums)
    m = max(ct.values())
    hold = [[] for _ in range(m)]
    for kk, v in ct.items():
        hold[v - 1].append(kk)
    return list(chain(*hold))[::-1][:k]

Quick select

def topKFrequent(self, nums: List[int], k: int) -> List[int]:
    count = collections.Counter(nums)
    unique = list(count.keys())
    
    def partition(left, right, pivot_index) -> int:
        pivot_frequency = count[unique[pivot_index]]
        unique[pivot_index], unique[right] = unique[right], unique[pivot_index]  
        
        store_index = left
        for i in range(left, right):
            if count[unique[i]] < pivot_frequency:
                unique[store_index], unique[i] = unique[i], unique[store_index]
                store_index += 1

        unique[right], unique[store_index] = unique[store_index], unique[right]  
        
        return store_index
    
    def quickselect(left, right, k_smallest) -> None:
        """
        Sort a list within left..right till kth less frequent element
        takes its place. 
        """
        # base case: the list contains only one element
        if left == right: 
            return
        
        # select a random pivot_index
        pivot_index = random.randint(left, right)     
                        
        # find the pivot position in a sorted list   
        pivot_index = partition(left, right, pivot_index)
        
        # if the pivot is in its final sorted position
        if k_smallest == pivot_index:
                return 
        # go left
        elif k_smallest < pivot_index:
            quickselect(left, pivot_index - 1, k_smallest)
        # go right
        else:
            quickselect(pivot_index + 1, right, k_smallest)
        
    n = len(unique) 
    # kth top frequent element is (n - k)th less frequent.
    # Do a partial sort: from less frequent to the most frequent, till
    # (n - k)th less frequent element takes its place (n - k) in a sorted array. 
    # All element on the left are less frequent.
    # All the elements on the right are more frequent.  
    quickselect(0, n - 1, n - k)
    # Return top k frequent elements
    return unique[n - k:]