Design Twitter

Difficulty: Medium Source: NeetCode

Problem

Design a simplified version of Twitter where users can post tweets, follow/unfollow another user, and see the 10 most recent tweets in their news feed.

Implement the Twitter class:

• Twitter() — initializes your Twitter object.
• postTweet(userId, tweetId) — composes a new tweet with ID tweetId by the user userId.
• getNewsFeed(userId) — retrieves the 10 most recent tweet IDs in the user’s news feed. Each item must be posted by users who the user followed or by the user themselves. Tweets are ordered from most recent to least recent.
• follow(followerId, followeeId) — the user with ID followerId starts following the user with ID followeeId.
• unfollow(followerId, followeeId) — the user with ID followerId stops following the user with ID followeeId.

Example 1: Input: postTweet(1, 5), getNewsFeed(1) → [5], follow(1, 2), postTweet(2, 6), getNewsFeed(1) → [6, 5], unfollow(1, 2), getNewsFeed(1) → [5]

Constraints:

• 1 <= userId, followerId, followeeId <= 500
• 0 <= tweetId <= 10^4
• All tweets have unique IDs
• At most 3 * 10^4 calls will be made to postTweet, getNewsFeed, follow, and unfollow

Prerequisites

Before attempting this problem, you should be comfortable with:

• Hash maps — storing per-user tweet lists and follow sets
• Min-heap merge — merging multiple sorted lists using a heap (k-way merge)
• Timestamps — using a monotonically increasing counter to order events

1. Brute Force — Collect and Sort

Intuition

For getNewsFeed, gather all tweets from the user and their followees into one big list, sort by timestamp descending, and return the top 10. Simple to implement but doesn’t scale when users have millions of tweets.

Algorithm

1. Store tweets as (timestamp, tweetId) per user in a list.
2. Store followees per user in a set (always include self).
3. For getNewsFeed: collect all tweets from user and followees, sort by timestamp descending, return first 10 tweet IDs.

Solution

from collections import defaultdict

class Twitter:
    def __init__(self):
        self.time = 0
        self.tweets = defaultdict(list)  # userId -> [(timestamp, tweetId)]
        self.following = defaultdict(set)  # followerId -> {followeeId}

    def postTweet(self, userId, tweetId):
        self.tweets[userId].append((self.time, tweetId))
        self.time += 1

    def getNewsFeed(self, userId):
        all_tweets = list(self.tweets[userId])
        for followee in self.following[userId]:
            all_tweets.extend(self.tweets[followee])
        all_tweets.sort(key=lambda x: -x[0])
        return [tid for _, tid in all_tweets[:10]]

    def follow(self, followerId, followeeId):
        self.following[followerId].add(followeeId)

    def unfollow(self, followerId, followeeId):
        self.following[followerId].discard(followeeId)


# Test
t = Twitter()
t.postTweet(1, 5)
print(t.getNewsFeed(1))   # [5]
t.follow(1, 2)
t.postTweet(2, 6)
print(t.getNewsFeed(1))   # [6, 5]
t.unfollow(1, 2)
print(t.getNewsFeed(1))   # [5]

Complexity

• Time: O(n log n) per getNewsFeed where n = total tweets by user + followees
• Space: O(n) total tweet storage

2. Min-Heap K-Way Merge

Intuition

Each user’s tweet list is already sorted by time (we append in order). getNewsFeed is really just merging the tweet lists of the user and their followees and taking the top 10. We can do this efficiently with a min-heap: seed the heap with the most recent tweet from each relevant user, then repeatedly extract the latest tweet and advance that user’s pointer — at most 10 times total.

We store (-timestamp, tweetId, userId, tweet_index) in the heap. The negative timestamp makes Python’s min-heap behave as a max-heap by recency.

Algorithm

1. Same data storage as brute force.
2. For getNewsFeed: a. Consider the user and all their followees. b. For each, if they have tweets, push (-timestamp, tweetId, userId, last_index) onto the heap. c. Pop up to 10 times: record the tweetId, then push the next tweet from the same user (if any) onto the heap.

Solution

import heapq
from collections import defaultdict

class Twitter:
    def __init__(self):
        self.time = 0
        self.tweets = defaultdict(list)  # userId -> [(timestamp, tweetId)]
        self.following = defaultdict(set)

    def postTweet(self, userId, tweetId):
        self.tweets[userId].append((self.time, tweetId))
        self.time += 1

    def getNewsFeed(self, userId):
        heap = []
        # Include the user themselves
        relevant = self.following[userId] | {userId}

        for uid in relevant:
            user_tweets = self.tweets[uid]
            if user_tweets:
                idx = len(user_tweets) - 1
                ts, tid = user_tweets[idx]
                heapq.heappush(heap, (-ts, tid, uid, idx - 1))

        feed = []
        while heap and len(feed) < 10:
            neg_ts, tid, uid, next_idx = heapq.heappop(heap)
            feed.append(tid)
            if next_idx >= 0:
                ts, next_tid = self.tweets[uid][next_idx]
                heapq.heappush(heap, (-ts, next_tid, uid, next_idx - 1))

        return feed

    def follow(self, followerId, followeeId):
        self.following[followerId].add(followeeId)

    def unfollow(self, followerId, followeeId):
        self.following[followerId].discard(followeeId)


# Test 1
t = Twitter()
t.postTweet(1, 5)
print(t.getNewsFeed(1))   # [5]
t.follow(1, 2)
t.postTweet(2, 6)
print(t.getNewsFeed(1))   # [6, 5]
t.unfollow(1, 2)
print(t.getNewsFeed(1))   # [5]

# Test 2 — self-tweets always in feed
t2 = Twitter()
t2.postTweet(1, 1)
t2.postTweet(1, 2)
t2.postTweet(2, 3)
t2.follow(1, 2)
print(t2.getNewsFeed(1))  # [3, 2, 1]

Complexity

• Time: O(f log f) per getNewsFeed where f = number of followees (heap seed), then O(10 log f) for extraction — effectively O(f log f)
• Space: O(f) for the heap, O(total tweets) for storage

Common Pitfalls

Not including the user in their own news feed. The problem says the feed includes the user’s own tweets. Don’t forget to add userId to the relevant set.

Using remove instead of discard for unfollow. If the user tries to unfollow someone they don’t follow, set.remove() raises a KeyError. Use set.discard() which silently does nothing.

Timestamp collisions. If you use real time (time.time()), tweets posted in the same millisecond can have equal timestamps. Use a simple integer counter that increments with every postTweet call.