My old solutions to the "Top K Frequent Elements" problem on LeetCode: https://leetcode.com/problems/top-k-frequent-elements/

topKFrequent.py

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        # sorting
        '''
        nums_map = collections.Counter(nums)
        return sorted(nums_map, key=nums_map.get)[-k:]
        # Equivalent to: heapq.nlargest(k, nums_map.keys(), key=count.get) 
        '''
        # heap way (equivalent to above)
        '''
        count_map = collections.Counter(nums)
        return heapq.nlargest(k, count_map.keys(), key=count_map.get)
        # Equivalent to: sorted(iterable, key=key, reverse=True)[:n]   
        '''
        
        # follow up: better than nlogn
        # Quickselect: ds used for kth something (smallest, largest, least freq, most freq ...) 
        # Avg time complexity: O(n)
        # worst case: O(n^2) BUT NEGLIGIBLE
        
        # when trying to go from nlogn to n
        # realize that we're doing more work than needed, by sorting we're getting the final
        # position for every eleemnt but we only need the last k
        
        # notice (N=len(nums)) N-k th position will give us the kth largest item in sorted order
        # when our final pivot position is N-k, then we have found the kth most frequent elements
        # because each partion will result in elements less than the pivot on the left
        # and elements greater than the pivot on the right and the final position of the pivot
        # will be its final position in the sorted array!
        # if pivot index > N-k then we recursively partion with a new pivot on the left side
        # if pivot index < N-k then we recursively partion with a new pivot on the right side
        
        if len(nums) == k:
            return nums
        
        freq_map = collections.Counter(nums)
        freq = list(freq_map.keys())
        
        def partition(left, right, pivot):
            pivot_freq = freq_map[freq[pivot]]
            
            # swap pivot element with last index
            freq[pivot], freq[right] = freq[right], freq[pivot]
            pivot = right
            
            boundary = left # boundary between elements less than and greater than pivot
            for i in range(left, right): # i from left to right-1
                if freq_map[freq[i]] < pivot_freq:
                    freq[i], freq[boundary] = freq[boundary], freq[i]
                    boundary += 1
            
            # swap back pivot into boundary
            # value in boundary should be in right side, so this swap puts it where it should
            # be, and now all elements less than and greater than pivot on either side
            freq[pivot], freq[boundary] = freq[boundary], freq[pivot]
            pivot = boundary
            return pivot
        # dont actually need to reassign pivot to point to its value
        # we could just save the pivot frequency before first swap (did this)
        # and then the second swap can be between right and boundary and we return boundary
        
        def quick_select(left=0, right=len(freq)-1):
            # using random pivot
            pivot = partition(left, right, random.randint(left, right))
            if pivot == len(freq)-k:
                # is the kth element in the correct position if the array was sorted
                return freq[-k:]
            elif pivot < len(freq)-k:
                # need to look for pivot that will become the kth sorted position
                # this in < kth position and we want to get to the kth position so check right
                return quick_select(pivot+1, right)
            else:
                return quick_select(left, pivot-1)
            
        return quick_select()