refactor: Improve time complexity using heap for topKFrequent

Author: KwonNayeon

top-k-frequent-elements/KwonNayeon.py

@@ -18,16 +18,43 @@
     - O(n)
 """
 
+# Original Solution
 class Solution:
     def topKFrequent(self, nums: List[int], k: int) -> List[int]:
         frequency = {}
         result = []
+        
         for num in nums:
             if num in frequency:
                 frequency[num] += 1
             else:
                 frequency[num] = 1
+                
         sorted_frequency = sorted(frequency.items(), key=lambda x: x[1], reverse=True)
+        
         for i in range(k):
             result.append(sorted_frequency[i][0])
+            
+        return result
+
+# Improved Solution using Heap
+# Time Complexity: O(n log k)
+class Solution:
+    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
+        frequency = {}
+        result = []
+        heap = []
+        
+        for num in nums:
+            if num in frequency:
+                frequency[num] += 1
+            else:
+                frequency[num] = 1
+                
+        for num, freq in frequency.items():
+            heapq.heappush(heap, (-freq, num))
+            
+        for i in range(k):
+            result.append(heapq.heappop(heap)[1])
+            
         return result