Read-only web browser

Solutions

Browse the interview solutions folder with syntax highlighting.

design-twitter.py

heap/design-twitter.py · Python · 3.1 KiB · 2026-03-21 14:56

Back to folder
# 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