692. Top K Frequent Words.py

class Solution:
    """
    {
    i:2
    love:2
    leetcode:1
    codong:1
    }
    怎麼處理k=1
    heqp可能會丟掉i而不是丟掉love
    這樣最終會回傳love而不是i
    same freq
    """
    def topKFrequent(self, words: List[str], k: int) -> List[str]:
        counter = collections.defaultdict(int)
        for word in words:
            counter[word]+=1
        heap = [] #T:O(k) and S:O(n)
        for word, freq in counter.items(): #O(nlogk)
            heapq.heappush(heap, (freq, word))
            if len(heap) > k:
                freq, word = heapq.heappop(heap)
                # print(word, freq, heap)
                if freq == heap[0][0]:
                    # same freq, we push it back
                    heapq.heappush(heap, (freq, word))
        return [i[1] for i in sorted(heap, key = lambda x: (-x[0], x[1]))][:k] # O(klogk)

class Solution:
    # Time Complexity = O(n + nlogk)
    # Space Complexity = O(n)
    def topKFrequent(self, words, k):
        count = collections.Counter(words)
        heap = []
        for key, value in count.items():
            heapq.heappush(heap, Word(value, key))
            if len(heap) > k:
                heapq.heappop(heap)
        print([(i.freq, i.word) for i in heap])
        res = []
        for _ in range(k):
            res.append(heapq.heappop(heap).word)
        return res[::-1]

class Word:
    def __init__(self, freq, word):
        self.freq = freq
        self.word = word
    
    def __lt__(self, other):
        """
        <
        """
        if self.freq == other.freq:
            return self.word > other.word
        return self.freq < other.freq
    
    def __eq__(self, other):
        """
        相等
        """
        return self.freq == other.freq and self.word == other.word