design-twitter.py
heap/design-twitter.py · Python · 3.1 KiB · 2026-03-21 14:56
# Sort, or min heap. Store tweet ids with timestamps in the heap, then use the merge k sorted lists approach. # heap import heapq from collections import defaultdict class Twitter: def __init__(self): self.tweetmap : dict[int, list[tuple[int, int]]] = defaultdict(list) self.followmap : dict[int, set] = defaultdict(set) self.time = 0 # O(1) time def postTweet(self, userId: int, tweetId: int) -> None: self.tweetmap[userId].append((-self.time, tweetId)) self.time += 1 # Heap, optimized - like merging k sorted lists # O(n log k) time, where n is the total number of tweets of all followers + user itself, and k is the number of followers + 1 (user itself) def getNewsFeed(self, userId: int) -> list[int]: res = [] heap = [] # Add to the heap only the most recent tweet of all followers + user itself feedsources = self.followmap[userId].copy() feedsources.add(userId) for feedsource in feedsources: if feedsource in self.tweetmap: index = len(self.tweetmap[feedsource]) - 1 time, tweet = self.tweetmap[feedsource][index] heapq.heappush(heap, [time, tweet, feedsource, index - 1]) # store the index of the next most recent tweet for this source while heap and len(res) < 10: time, tweet, feedsource, index = heapq.heappop(heap) res.append(tweet) # if the feed source still has tweets, add the next more recent if index >= 0: time, tweet = self.tweetmap[feedsource][index] heapq.heappush(heap, [time, tweet, feedsource, index - 1]) return res # Heap, not optimized (insert all tweets of all followers) # O(n log n) time, where n is the total number of tweets of all followers + user itself def getNewsFeed1(self, userId: int) -> list[int]: tweets = [] for followee in self.followmap[userId]: for tw in self.tweetmap[followee]: heapq.heappush(tweets, tw) for tw in self.tweetmap[userId]: heapq.heappush(tweets, tw) res = [] for tw in range(min(10, len(tweets))): res.append(heapq.heappop(tweets)[1]) return res # Sorting - O(n*m * t log t) time, where n is the number of followers + 1 (user itself), m is the number of tweets per user, and t is the total number of tweets def getNewsFeed2(self, userId: int) -> list[int]: feed = self.tweetmap[userId][:] # make a copy for followee in self.followmap[userId]: feed.extend(self.tweetmap[followee]) feed.sort() return [tw for _, tw in feed[:10]] # O(1) time def follow(self, followerId: int, followeeId: int) -> None: if followeeId != followerId: self.followmap[followerId].add(followeeId) # O(1) time def unfollow(self, followerId: int, followeeId: int) -> None: self.followmap[followerId].discard(followeeId) # like .remove() but doesn't raise exc if the element is not there