Design Twitter

Design Twitter

Design Twitter

Problem

Design a simple version of Twitter. Users can post tweets, follow or unfollow other users, and see the 10 most recent tweets in their news feed.

Implement the Twitter class:

• Twitter(): Create the object.
• postTweet(userId, tweetId): User userId posts a new tweet with ID tweetId. Every tweetId is unique.
• getNewsFeed(userId): Return the 10 most recent tweet IDs in the user's feed. The tweets must be from the user and from the users they follow. Order them from most recent to least recent.
• follow(followerId, followeeId): followerId starts following followeeId.
• unfollow(followerId, followeeId): followerId stops following followeeId.

Examples

Example 1

Input
["Twitter", "postTweet", "getNewsFeed", "follow", "postTweet", "getNewsFeed", "unfollow", "getNewsFeed"]
[[], [1, 5], [1], [1, 2], [2, 6], [1], [1, 2], [1]]

Output
[null, null, [5], null, null, [6, 5], null, [5]]

• User 1 posts tweet 5. The feed of user 1 is [5].
• User 1 follows user 2. User 2 posts tweet 6.
• The feed of user 1 is [6, 5]. Tweet 6 is newer, so it comes first.
• User 1 unfollows user 2. The feed is [5] again.

Constraints

• 1 <= userId, followerId, followeeId <= 500
• 0 <= tweetId <= 10^4
• All tweet IDs are unique.
• At most 3 * 10^4 calls are made.
• A user cannot follow themself.

Main Idea

We need to answer three questions fast:

• Who does a user follow? We use a Map from userId to a Set of followee IDs. A Set makes follow and unfollow O(1).
• What did a user tweet? We use a Map from userId to an array of [time, tweetId]. We add new tweets at the end, so each array is sorted from old to new.
• Which tweet is newest? We use a global time counter. It goes up by 1 on every post. A bigger time means a newer tweet.

To build the feed, we need the tweets with the biggest time. A max heap (sorted by time) gives us the newest tweet first.

The Heap Class Used in the Solution

The solution uses a heap class with a custom comparator. It is a min heap by default. For a max heap by time, we pass (a, b) => b[0] - a[0].

class newHeap {
    constructor(comparator = (a, b) => a - b) {
        this.heap = [];
        this.comparator = comparator;
    }

    size() {
        return this.heap.length;
    }

    isEmpty() {
        return this.heap.length === 0;
    }

    peek() {
        return this.heap[0];
    }

    offer(value) {
        this.heap.push(value);
        this.bubbleUp(this.heap.length - 1);
    }

    poll() {
        if (this.isEmpty()) return null;

        if (this.heap.length === 1) {
            return this.heap.pop();
        }

        let result = this.heap[0];

        this.heap[0] = this.heap.pop();
        this.bubbleDown(0);

        return result;
    }

    bubbleUp(index) {
        while (index > 0) {
            let parent = Math.floor((index - 1) / 2);

            if (this.comparator(this.heap[index], this.heap[parent]) >= 0) {
                break;
            }

            this.swap(index, parent);
            index = parent;
        }
    }

    bubbleDown(index) {
        while (true) {
            let left = index * 2 + 1;
            let right = index * 2 + 2;
            let best = index;

            if (
                left < this.heap.length &&
                this.comparator(this.heap[left], this.heap[best]) < 0
            ) {
                best = left;
            }

            if (
                right < this.heap.length &&
                this.comparator(this.heap[right], this.heap[best]) < 0
            ) {
                best = right;
            }

            if (best === index) {
                break;
            }

            this.swap(index, best);
            index = best;
        }
    }

    swap(i, j) {
        [this.heap[i], this.heap[j]] = [this.heap[j], this.heap[i]];
    }
}

How the Heap Works

• offer adds a value at the end, then moves it up until the heap rule is true.
• poll removes the top value. It moves the last value to the top, then moves it down until the heap rule is true.
• The comparator decides who goes on top. Here, b[0] - a[0] puts the biggest time on top.
• offer and poll both take O(log n). peek takes O(1).

Solution 1: Put All Tweets in One Heap

We collect every tweet of the user and of all followees into one heap. Then we take the top 10.

class Twitter {
    constructor() {
        this.time = 0;
        this.following = new Map(); // userId -> Set of followee IDs
        this.tweet = new Map();     // userId -> array of [time, tweetId]
    }

    postTweet(userId, tweetId) {
        if (!this.tweet.has(userId)) {
            this.tweet.set(userId, []);
        }

        this.tweet.get(userId).push([this.time, tweetId]);
        this.time++;
    }

    getNewsFeed(userId) {
        let followingList;

        if (this.following.has(userId)) {
            followingList = this.following.get(userId);
        } else {
            followingList = new Set();
        }

        // the user also sees their own tweets
        followingList.add(userId);

        // max heap by time: the newest tweet is on top
        let heap = new newHeap((a, b) => b[0] - a[0]);
        let ans = [];

        for (let follower of followingList) {
            if (this.tweet.has(follower)) {
                let posts = this.tweet.get(follower);

                for (let post of posts) {
                    heap.offer(post);
                }
            }
        }

        // take the 10 newest tweets
        while (!heap.isEmpty() && ans.length != 10) {
            ans.push(heap.poll()[1]);
        }

        return ans;
    }

    follow(followerId, followeeId) {
        if (!this.following.has(followerId)) {
            this.following.set(followerId, new Set());
        }

        this.following.get(followerId).add(followeeId);
    }

    unfollow(followerId, followeeId) {
        if (
            this.following.has(followerId) &&
            this.following.get(followerId).has(followeeId)
        ) {
            this.following.get(followerId).delete(followeeId);
        }
    }
}

Tip: In getNewsFeed, the line followingList.add(userId) changes the saved set, so the user is added to their own follow list. It does not break the answers here. But it is cleaner to make a copy: new Set(this.following.get(userId) || []), and then add userId to the copy.

How the Code Works

The constructor

• time is a counter that gives each tweet a time stamp.
• following stores who each user follows.
• tweet stores the tweets of each user.

postTweet

• If the user has no tweet list yet, we create one.
• We push [time, tweetId] and increase time.

getNewsFeed

• Get the follow list of the user. If there is none, use an empty set.
• Add the user to the list, so their own tweets are included.
• Put all tweets of everyone in the list into a max heap.
• Poll the heap up to 10 times. Each poll gives the newest remaining tweet. We save only [1], the tweetId.

follow and unfollow

• follow creates the set if needed, and adds the followee.
• unfollow removes the followee, only if it is there.

Dry Run

Let us run the example.

Complexity

Let T be the total number of tweets by the user and all followees.

• postTweet, follow, unfollow: O(1)
• getNewsFeed: O(T log T)
  • We push T tweets into the heap. Each push takes O(log T).
  • The 10 polls take O(10 log T), which is small.
• Space Complexity: O(T) for the heap in getNewsFeed.
  • The saved data (following and tweet) takes O(U + F + P), where U is the number of users, F is the number of follow links, and P is the number of posted tweets.

This is fine for the constraints. But we push many tweets that we never use. The next solution fixes this.

Solution 2: Heap with One Tweet per User (Optimized)

We only need the 10 newest tweets. So we do not need to put every tweet in the heap.

Each user's tweet list is already sorted by time. So the newest tweet of a user is the last one in their list. We use this like merging sorted lists:

• Put only the newest tweet of each followee in the heap.
• Poll the top. It is the newest tweet of all.
• After polling a tweet from user f, put the next older tweet of f into the heap.
• Repeat until we have 10 tweets or the heap is empty.

Each heap item is [time, tweetId, userId, index]. The userId and index tell us where to find the next older tweet.

class Twitter {
    constructor() {
        this.time = 0;
        this.following = new Map();
        this.tweet = new Map();
    }

    postTweet(userId, tweetId) {
        if (!this.tweet.has(userId)) {
            this.tweet.set(userId, []);
        }

        this.tweet.get(userId).push([this.time, tweetId]);
        this.time++;
    }

    getNewsFeed(userId) {
        // copy the follow list, so the saved set does not change
        let followees = new Set(this.following.get(userId) || []);
        followees.add(userId);

        let heap = new newHeap((a, b) => b[0] - a[0]);

        // put only the newest tweet of each user
        for (let f of followees) {
            let posts = this.tweet.get(f);

            if (posts && posts.length > 0) {
                let idx = posts.length - 1;
                heap.offer([posts[idx][0], posts[idx][1], f, idx]);
            }
        }

        let ans = [];

        while (!heap.isEmpty() && ans.length < 10) {
            let [time, tweetId, f, idx] = heap.poll();
            ans.push(tweetId);

            // put the next older tweet of the same user
            if (idx > 0) {
                let prev = this.tweet.get(f)[idx - 1];
                heap.offer([prev[0], prev[1], f, idx - 1]);
            }
        }

        return ans;
    }

    follow(followerId, followeeId) {
        if (!this.following.has(followerId)) {
            this.following.set(followerId, new Set());
        }

        this.following.get(followerId).add(followeeId);
    }

    unfollow(followerId, followeeId) {
        if (this.following.has(followerId)) {
            this.following.get(followerId).delete(followeeId);
        }
    }
}

How getNewsFeed Works Here

• We copy the follow set and add the user.
• The heap starts with at most one tweet per followee. So its size is at most F, the number of followees.
• Each poll gives the newest tweet. Then we add the next older tweet from the same user, because that user's older tweet may now be the next newest.
• We stop after 10 tweets.

Dry Run

Let us take user 1 who follows user 2. User 1 has tweets [[0, 5], [2, 7]]. User 2 has tweets [[1, 6]].

The feed is [7, 6, 5].

Complexity

Let F be the number of followees (plus the user).

• postTweet, follow, unfollow: O(1)
• getNewsFeed: O(F log F)
  • We push at most F items first.
  • Then we do at most 10 polls and 10 pushes. Each takes O(log F).
• Space Complexity: O(F) for the heap in getNewsFeed.
  • The saved data takes O(U + F_total + P), as before.

Comparison

Summary

• We store who follows whom in a Map of Sets, and tweets of each user in a Map of arrays.
• A global time counter tells us which tweet is newer.
• A max heap by time gives the newest tweet first.
• Solution 1 puts every tweet in the heap. It is simple, but it does extra work.
• Solution 2 puts only the newest tweet of each user in the heap. After each poll, it adds the next older tweet of that user. This works like merging sorted lists.
• Make a copy of the follow set in getNewsFeed, so you do not change the saved data.
• postTweet, follow and unfollow take O(1). getNewsFeed takes O(T log T) in Solution 1 and O(F log F) in Solution 2.