MUYANG GUO / INDEX

LeetCode

LeetCode 347 Top K Frequent Elements - Medium

347. Top K Frequent Elements -- Medium

·2 min read·#LeetCode#Medium#Python

347. Top K Frequent Elements — Medium

Open on LeetCode

Problem

  1. Top K Frequent Elements -- Medium

Given a non-empty array of integers, return the k most frequent elements.

Example 1: Input: nums = [1,1,1,2,2,3], k = 2 Output: [1,2]

Example 2: Input: nums = [1], k = 1 Output: [1]

Note: You may assume k is always valid, 1 ≤ k ≤ number of unique elements. Your algorithm's time complexity must be better than O(n log n), where n is the array's size. It's guaranteed that the answer is unique, in other words the set of the top k frequent elements is unique. You can return the answer in any order.

Solution

### O(klogN), with heap:
import heapq
class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        map = dict()
        for num in nums:
            map[num] = map.get(num, 0) + 1
        heap = []
        for key, freq in map.items():
            heapq.heappush(heap, (-freq, key))
        res = []
        for i in range(k):
            _, key = heapq.heappop(heap)
            res.append(key)
        
        return res
 
### Qucikselect, Average O(N), Worst Case O(N^2):
from collections import Counter
class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        count = Counter(nums)
        unique = list(count.keys())
        
        def partition(left, right, pivot_index) -> int:
            pivot_frequency = count[unique[pivot_index]]
            # 1. move pivot to end
            unique[pivot_index], unique[right] = unique[right], unique[pivot_index]  
            
            # 2. move all less frequent elements to the left
            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
 
            # 3. move pivot to its final place
            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:]

Comments