diff --git a/trees_and_graphs/top_k_frequent_heap.py b/trees_and_graphs/top_k_frequent_heap.py new file mode 100644 index 0000000..8acc9d1 --- /dev/null +++ b/trees_and_graphs/top_k_frequent_heap.py @@ -0,0 +1,17 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +# author: bt3gl + +def top_k_frequentnums: list[int], k: int) -> list[int]: + + # O(1) time + if k == len(nums): + return nums + + # 1. build a hashmap element: frequency + counter = Counter(nums) + + # 2. build a heap of k most frequent elements + # 3. build an output array + # O(N log k) time + return heapq.nlargest(k, counter.keys(), key=counter.get)