Design Twitter

mediumAmazon-tagged: Design-style coding63 stepsJava · Python · C++

Post tweets, follow and unfollow, and get the 10 most recent tweets from yourself and the people you follow. Each user keeps their own tweets, newest last, with a global timestamp. The news feed merges those lists with a max-heap by time and stops after 10.

▶ Watch the animated trace

Traced input: postTweet(1, 5), getNewsFeed(1), follow(1, 2), postTweet(2, 6), getNewsFeed(1), unfollow(1, 2), getNewsFeed(1) → [null, null, [5], null, null, [6, 5], null, [5]]. Original problem: neetcode.io/problems/design-twitter-feed

Complexity

BestO(1)
AverageO(F + 10 log F)
WorstO(F + 10 log F)
SpaceO(users + tweets)

F = number of followees. post, follow and unfollow are O(1).

Key idea

Common mistakes

Solution

Java
class Twitter {
    private int time = 0;
    private final Map<Integer, List<int[]>> tweets = new HashMap<>();
    private final Map<Integer, Set<Integer>> follows = new HashMap<>();

    void postTweet(int userId, int tweetId) {
        tweets.computeIfAbsent(userId, k -> new ArrayList<>()).add(new int[]{time++, tweetId});
    }

    List<Integer> getNewsFeed(int userId) {
        PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) -> b[0] - a[0]);
        Set<Integer> users = new HashSet<>(follows.getOrDefault(userId, Set.of()));
        users.add(userId);
        for (int u : users) {
            List<int[]> list = tweets.get(u);
            if (list != null && !list.isEmpty()) pq.offer(new int[]{list.get(list.size() - 1)[0], list.get(list.size() - 1)[1], u, list.size() - 1});
        }
        List<Integer> feed = new ArrayList<>();
        while (!pq.isEmpty() && feed.size() < 10) {
            int[] top = pq.poll();
            feed.add(top[1]);
            if (top[3] > 0) {
                int[] prev = tweets.get(top[2]).get(top[3] - 1);
                pq.offer(new int[]{prev[0], prev[1], top[2], top[3] - 1});
            }
        }
        return feed;
    }

    void follow(int a, int b) {
        if (a != b) follows.computeIfAbsent(a, k -> new HashSet<>()).add(b);
    }

    void unfollow(int a, int b) {
        if (follows.containsKey(a)) follows.get(a).remove(b);
    }
}
Python
class Twitter:
    def __init__(self):
        self.time = 0
        self.tweets = defaultdict(list)
        self.follows = defaultdict(set)

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

    def getNewsFeed(self, userId: int) -> List[int]:
        heap = []
        for u in self.follows[userId] | {userId}:
            if self.tweets[u]:
                i = len(self.tweets[u]) - 1
                t, tid = self.tweets[u][i]
                heap.append((-t, tid, u, i))
        heapq.heapify(heap)
        feed = []
        while heap and len(feed) < 10:
            _, tid, u, i = heapq.heappop(heap)
            feed.append(tid)
            if i > 0:
                t, prev = self.tweets[u][i - 1]
                heapq.heappush(heap, (-t, prev, u, i - 1))
        return feed

    def follow(self, followerId: int, followeeId: int) -> None:
        if followerId != followeeId: self.follows[followerId].add(followeeId)

    def unfollow(self, followerId: int, followeeId: int) -> None:
        self.follows[followerId].discard(followeeId)
C++
class Twitter {
    int time = 0;
    unordered_map<int, vector<pair<int, int>>> tweets;
    unordered_map<int, unordered_set<int>> follows;

public:
    void postTweet(int userId, int tweetId) {
        tweets[userId].push_back({time++, tweetId});
    }

    vector<int> getNewsFeed(int userId) {
        priority_queue<array<int, 4>> pq;  // {time, tweetId, user, index}
        unordered_set<int> users = follows[userId];
        users.insert(userId);
        for (int u : users) {
            auto& list = tweets[u];
            if (!list.empty()) pq.push({list.back().first, list.back().second, u, (int)list.size() - 1});
        }
        vector<int> feed;
        while (!pq.empty() && feed.size() < 10) {
            auto [t, id, u, i] = pq.top(); pq.pop();
            feed.push_back(id);
            if (i > 0) {
                auto& prev = tweets[u][i - 1];
                pq.push({prev.first, prev.second, u, i - 1});
            }
        }
        return feed;
    }

    void follow(int followerId, int followeeId) {
        if (followerId != followeeId) follows[followerId].insert(followeeId);
    }

    void unfollow(int followerId, int followeeId) {
        follows[followerId].erase(followeeId);
    }
};

Edge cases

Each one has its own trace in the interactive view.

CaseInputAnswer
Unfollow yourself
You always see your own tweets.
postTweet(1, 7), unfollow(1, 1), getNewsFeed(1)[null, null, null, [7]]
More than 10 tweets
Only the 10 most recent.
postTweet(1, 0) … postTweet(1, 11), getNewsFeed(1)[11, 10, 9, 8, 7, 6, 5, 4, 3, 2]
← Snapshot ArrayMaximum Frequency Stack →