Design Twitter

Sanjeev SharmaSanjeev Sharma
3 min read

Advertisement

Problem

Implement Twitter: postTweet, getNewsFeed (top 10 most recent from user + followees), follow, unfollow.

Key insight: K-way merge on followee tweet lists (sorted by timestamp descending). Min-heap of size 10 or max-heap merge.

Solutions

# Python
import heapq
from collections import defaultdict

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

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

    def getNewsFeed(self, userId):
        heap = []  # (-time, tweetId, userId, index)
        users = self.follows[userId] | {userId}
        for uid in users:
            tweets = self.tweets[uid]
            if tweets:
                i = len(tweets) - 1
                t, tid = tweets[i]
                heapq.heappush(heap, (-t, tid, uid, i - 1))
        result = []
        while heap and len(result) < 10:
            neg_t, tid, uid, idx = heapq.heappop(heap)
            result.append(tid)
            if idx >= 0:
                t2, tid2 = self.tweets[uid][idx]
                heapq.heappush(heap, (-t2, tid2, uid, idx - 1))
        return result

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

    def unfollow(self, followerId, followeeId):
        self.follows[followerId].discard(followeeId)
// C++
class Twitter {
    int timer = 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({timer++, tweetId});
    }
    vector<int> getNewsFeed(int userId) {
        using T = tuple<int,int,int,int>; // time, tweetId, userId, idx
        priority_queue<T> maxH;
        follows[userId].insert(userId);
        for (int uid : follows[userId]) {
            if (!tweets[uid].empty()) {
                int i = tweets[uid].size()-1;
                maxH.push({tweets[uid][i].first, tweets[uid][i].second, uid, i-1});
            }
        }
        vector<int> res;
        while (!maxH.empty() && (int)res.size() < 10) {
            auto [t, tid, uid, idx] = maxH.top(); maxH.pop();
            res.push_back(tid);
            if (idx >= 0) maxH.push({tweets[uid][idx].first, tweets[uid][idx].second, uid, idx-1});
        }
        return res;
    }
    void follow(int a, int b) { follows[a].insert(b); }
    void unfollow(int a, int b) { follows[a].erase(b); }
};
// Java
class Twitter {
    int timer = 0;
    Map<Integer, List<int[]>> tweets = new HashMap<>();
    Map<Integer, Set<Integer>> follows = new HashMap<>();
    public void postTweet(int userId, int tweetId) {
        tweets.computeIfAbsent(userId, k->new ArrayList<>()).add(new int[]{timer++, tweetId});
    }
    public List<Integer> getNewsFeed(int userId) {
        PriorityQueue<int[]> maxH = new PriorityQueue<>((a,b)->b[0]-a[0]);
        Set<Integer> users = follows.getOrDefault(userId, new HashSet<>());
        users = new HashSet<>(users); users.add(userId);
        for (int uid : users) {
            List<int[]> tw = tweets.getOrDefault(uid, new ArrayList<>());
            if (!tw.isEmpty()) { int i=tw.size()-1; maxH.offer(new int[]{tw.get(i)[0],tw.get(i)[1],uid,i-1}); }
        }
        List<Integer> res = new ArrayList<>();
        while (!maxH.isEmpty() && res.size() < 10) {
            int[] p = maxH.poll(); res.add(p[1]);
            if (p[3] >= 0) { int[]tw2=tweets.get(p[2]).get(p[3]); maxH.offer(new int[]{tw2[0],tw2[1],p[2],p[3]-1}); }
        }
        return res;
    }
    public void follow(int a, int b) { follows.computeIfAbsent(a,k->new HashSet<>()).add(b); }
    public void unfollow(int a, int b) { if(follows.containsKey(a))follows.get(a).remove(b); }
}
// JavaScript
class Twitter {
    constructor() { this.tweets = new Map(); this.follows = new Map(); this.time = 0; }
    postTweet(userId, tweetId) {
        if (!this.tweets.has(userId)) this.tweets.set(userId, []);
        this.tweets.get(userId).push([this.time++, tweetId]);
    }
    getNewsFeed(userId) {
        const users = new Set([userId, ...(this.follows.get(userId) || [])]);
        const all = [];
        for (const uid of users) for (const t of (this.tweets.get(uid) || [])) all.push(t);
        return all.sort((a,b) => b[0]-a[0]).slice(0,10).map(t=>t[1]);
    }
    follow(a, b) { if (!this.follows.has(a)) this.follows.set(a, new Set()); this.follows.get(a).add(b); }
    unfollow(a, b) { if (this.follows.has(a)) this.follows.get(a).delete(b); }
}

Complexity

  • getNewsFeed: O(U * T log U) where U=followees, T=tweets per user (heap-based) or O(N log N)
  • Space: O(total tweets)

Advertisement

Sanjeev Sharma

Written by

Sanjeev Sharma

Full Stack Engineer · E-mopro