Amazon Online Assessment (OA) 2021 - Top K Frequently Mentioned Keywords | SHL

Practice on leetcode here: https://leetcode.com/problems/top-k-frequent-words/

1from heapq import heappop, heappush
2import re
3from typing import Counter, List
4
5class Down:
6    def __init__(self, value):
7        self.value = value
8
9    def __lt__(self, other):
10        return self.value > other.value
11
12def top_mentioned(k: int, keywords: List[str], reviews: List[str]) -> List[str]:
13    patt = re.compile(r'\b(:?{})\b'.format('|'.join(keywords)), flags=re.IGNORECASE)
14    counts = Counter(
15        word
16        for review in reviews
17        for word in set(match[0].lower() for match in patt.finditer(review))
18    )
19    queue = []
20    for word, count in counts.items():
21        heappush(queue, (count, Down(word)))
22        if len(queue) > k:
23            heappop(queue)
24    res = []
25    while len(queue) > 0:
26        res.append(heappop(queue)[1].value)
27    return res[::-1]
28
29if __name__ == '__main__':
30    k = int(input())
31    keywords = [input() for _ in range(int(input()))]
32    reviews = [input() for _ in range(int(input()))]
33    res = top_mentioned(k, keywords, reviews)
34    print(' '.join(res))
35