Start learning free

LESSON 404 OF 569 · HEAP

Design Twitter in Java: Merging Feeds with a Max-Heap

Practice O(F log F) Read the full explanation

Design a small version of Twitter. Users post tweets and follow each other, and each user has a news feed: the 10 most recent tweets posted by the user and by the people the user follows, newest first.

Implement the class Twitter:

  • Twitter() creates an empty Twitter: no users, no tweets.
  • postTweet(userId, tweetId) posts a new tweet with the id tweetId by the user userId. Each call posts later than all the calls before it.
  • getNewsFeed(userId) returns the ids of the 10 most recent tweets in the user's news feed, newest first. A feed with fewer than 10 tweets returns all of them, and an empty feed returns no ids at all.
  • follow(followerId, followeeId): the user followerId starts following the user followeeId.
  • unfollow(followerId, followeeId): the user followerId stops following the user followeeId.

Following is one way: when user 1 follows user 2, user 1 sees user 2's tweets, old and new, but not the other way round. Following someone you already follow, following yourself and unfollowing someone you don't follow change nothing; your own tweets are always in your feed.

The input is a list of operations, applied in order to one new Twitter. The output has one line for each getNewsFeed call: the tweet ids it returns, separated by spaces, or an empty line for an empty feed. The starter code spells these names the way your language does.

 

Example 1:

Input:
    postTweet(1, 5)
    getNewsFeed(1)
    follow(1, 2)
    postTweet(2, 6)
    getNewsFeed(1)
    unfollow(1, 2)
    getNewsFeed(1)
Output:
5
6 5
5
Explanation:
User 1's first feed holds their own tweet 5.
User 1 then follows user 2, who posts tweet 6: it is newer than 5, so it comes first.
After the unfollow, user 2's tweets leave user 1's feed again.

Example 2:

Input:
    follow(1, 2)
    postTweet(1, 101)
    postTweet(2, 201)
    postTweet(1, 102)
    postTweet(2, 202)
    ...
    postTweet(1, 106)
    postTweet(2, 206)
    getNewsFeed(1)
    getNewsFeed(2)
Output:
206 106 205 105 204 104 203 103 202 102
206 205 204 203 202 201
Explanation:
Users 1 and 2 take turns posting 6 tweets each (101 to 106 and 201 to 206).
User 1's feed has all 12, but shows only the 10 newest: 101 and 201 are left out.
User 2 doesn't follow user 1, so user 2's feed holds only user 2's own tweets.

Example 3:

Input:
    getNewsFeed(3)
    follow(3, 3)
    postTweet(3, 7)
    getNewsFeed(3)
    unfollow(3, 4)
    follow(3, 4)
    getNewsFeed(3)
    postTweet(4, 8)
    getNewsFeed(3)
Output:

7
7
8 7
Explanation:
The first line of the output is empty: nobody has posted yet, so user 3's feed is empty.
follow(3, 3) changes nothing: tweet 7 appears once. unfollow(3, 4) changes nothing either: user 3 didn't follow user 4.
User 3 follows user 4 before user 4 posts anything; once tweet 8 arrives, it is in user 3's feed.

 

Constraints:

  • 1 <= userId, followerId, followeeId <= 500
  • 0 <= tweetId <= 109
  • All the tweets have different ids.
  • At most 105 calls in total are made to postTweet, getNewsFeed, follow and unfollow.

Note:

postTweet, follow and unfollow should run in O(1) time, and getNewsFeed in O(F log F) time, where F is the number of people the user follows. Collecting and sorting all the tweets of those people at every getNewsFeed would take time that grows with every tweet ever posted.

STEP 1 OF 24
    Ctrl + Enter to check

    JAVA · HEAP

    Design Twitter in Java, explained

    Problem overview

    The class Twitter models users who post tweets and follow each other. It has four methods. postTweet(userId, tweetId) posts a tweet. getNewsFeed(userId) returns a List<Integer> with the ids of the 10 most recent tweets posted by the user and by the people the user follows, newest first. follow(followerId, followeeId) and unfollow(followerId, followeeId) change who follows whom.

    Several rules decide the special cases. Following is one way: if user 1 follows user 2, user 1 sees user 2's tweets, old and new, but not the other way round. A feed with fewer than 10 tweets returns all of them, and an empty feed returns an empty list. Following someone you already follow, following yourself and unfollowing someone you don't follow change nothing. Your own tweets are always in your feed. Tweet ids are not ordered by age: each call to postTweet is later than all the calls before it, whatever the id.

    Here is a short session, with the feed after each step.

    Twitter tw = new Twitter();
    tw.postTweet(1, 5);
    System.out.println(tw.getNewsFeed(1));   // user 1's own tweet
    tw.follow(1, 2);
    tw.postTweet(2, 6);
    System.out.println(tw.getNewsFeed(1));   // user 2's newer tweet comes first
    tw.unfollow(1, 2);
    System.out.println(tw.getNewsFeed(1));   // user 2's tweets are gone again
    
    [5]
    [6, 5]
    [5]
    

    Approach: a clock, per-user lists and a heap

    The simple approach keeps every tweet in one global list. getNewsFeed then walks the whole list, keeps the tweets whose author is in the user's group, and reads the last 10. That costs time proportional to every tweet ever posted, and it grows with each post. Collecting the tweets of the followed people and sorting them has the same problem.

    The lesson stores the data so that a feed only looks at a handful of tweets:

    • A clock, the int time, goes up by one with every tweet. Tweet ids cannot tell which tweet is newer, since tweet 106 may be newer than tweet 201, but the clock can.
    • A map tweets from each user id to that user's tweets, each stored as an int[] holding {time, tweetId}, oldest first. Because the clock only goes up, each new tweet belongs at the end of the list.
    • A map following from each user id to the set of ids that user follows. A Set ignores a second follow of the same person, and adding or removing takes O(1).

    A feed is then a merge of sorted lists. Each person's list is sorted by time, and the feed needs the 10 newest tweets across all those lists. This is the same idea as merging sorted lists with a heap: put the newest tweet of each person in a max-heap ordered by time, take the top, and replace it with the same person's previous tweet. After at most 10 such steps the feed is complete. Older tweets of a person are never looked at before their newer ones have been taken, so they can wait in their list.

    Step by step

    1. Declare the fields: the clock time starting at 0, the map tweets and the map following.
    2. postTweet: increase time first, so every tweet has a unique, bigger time than the ones before it.
    3. If the user has no list yet, put a new ArrayList<>() in tweets. A user gets a list only on the first post.
    4. Add new int[]{time, tweetId} to that list.
    5. getNewsFeed: build a list users that starts with userId, because your own tweets are always in your feed, even though you cannot follow yourself.
    6. If the user follows anyone, add all those ids with addAll. follow never stores a self-follow, so nobody appears twice.
    7. Create the heap newest with the comparator (a, b) -> b[0] - a[0], which puts the biggest time, the newest tweet, on top. Its entries are {time, user, index}: the tweet's time, its author, and its position in the author's list.
    8. For each person who has posted, compute index, the last position in their list, and add their newest tweet as an entry.
    9. Create the empty list feed, then loop while the heap is not empty and the feed has fewer than 10 ids.
    10. In the loop, poll() the top entry, read its user and index, and add that tweet's id to feed. Tweets come out newest first, so the ids land in the right order.
    11. If that person has an older tweet (index > 0), push it into the heap in the same way. The heap now holds that person's next candidate.
    12. Return feed.
    13. follow: return immediately if both ids are equal. Otherwise create the follower's set if it does not exist, and add the followee.
    14. unfollow: if the follower has a set, remove the followee from it. remove does nothing for a value that is not in the set.

    The solution in Java

    class Twitter {
        int time = 0;
        Map<Integer, List<int[]>> tweets = new HashMap<>();
        Map<Integer, Set<Integer>> following = new HashMap<>();
    
        public void postTweet(int userId, int tweetId) {
            time++;
            if (!tweets.containsKey(userId)) {
                tweets.put(userId, new ArrayList<>());
            }
            tweets.get(userId).add(new int[]{time, tweetId});
        }
    
        public List<Integer> getNewsFeed(int userId) {
            List<Integer> users = new ArrayList<>();
            users.add(userId);
            if (following.containsKey(userId)) {
                users.addAll(following.get(userId));
            }
            PriorityQueue<int[]> newest = new PriorityQueue<>((a, b) -> b[0] - a[0]);
            for (int user : users) {
                if (tweets.containsKey(user)) {
                    int index = tweets.get(user).size() - 1;
                    newest.add(new int[]{tweets.get(user).get(index)[0], user, index});
                }
            }
            List<Integer> feed = new ArrayList<>();
            while (!newest.isEmpty() && feed.size() < 10) {
                int[] top = newest.poll();
                int user = top[1], index = top[2];
                feed.add(tweets.get(user).get(index)[1]);
                if (index > 0) {
                    newest.add(new int[]{tweets.get(user).get(index - 1)[0], user, index - 1});
                }
            }
            return feed;
        }
    
        public void follow(int followerId, int followeeId) {
            if (followerId == followeeId) {
                return;
            }
            if (!following.containsKey(followerId)) {
                following.put(followerId, new HashSet<>());
            }
            following.get(followerId).add(followeeId);
        }
    
        public void unfollow(int followerId, int followeeId) {
            if (following.containsKey(followerId)) {
                following.get(followerId).remove(followeeId);
            }
        }
    }
    

    A few details about the Java code:

    • The comparator. b[0] - a[0] sorts by the first number, largest first. Subtraction can overflow for large values, but times are small here (at most 100,000 tweets), so it is safe. For arbitrary ints, Integer.compare(b[0], a[0]) is the safe form.
    • int[] as a record. The entries are plain arrays: {time, tweetId} in the lists and {time, user, index} in the heap. The positions mean different things in each, so keep them straight.
    • Auto-unboxing. for (int user : users) unboxes each Integer from the list. following.get(followerId).remove(followeeId) works on a Set<Integer> and removes the value, because a Set has no remove(int index). On a List<Integer> the same call would remove by position, a classic trap.
    • The empty feed. If nobody in the group has posted, the heap stays empty, the loop never runs and the method returns an empty ArrayList, not null.

    Seeing it work

    When two users take turns posting six tweets each (ids 101 to 106 for user 1 and 201 to 206 for user 2), and user 1 follows user 2, the feeds are:

    Twitter t2 = new Twitter();
    t2.follow(1, 2);
    for (int i = 1; i <= 6; i++) {
        t2.postTweet(1, 100 + i);
        t2.postTweet(2, 200 + i);
    }
    System.out.println(t2.getNewsFeed(1));
    System.out.println(t2.getNewsFeed(2));
    
    [206, 106, 205, 105, 204, 104, 203, 103, 202, 102]
    [206, 205, 204, 203, 202, 201]
    

    User 1 has 12 tweets in the feed but sees only the 10 newest, so 101 and 201 are left out. User 2 does not follow user 1, so user 2 sees only their own six.

    The special cases behave as the rules say. A new Twitter returns an empty list for any user. After follow(3, 3) and postTweet(3, 7), the feed of user 3 shows tweet 7 once, and the map following is still empty, since a self-follow is never stored. After unfollow(3, 4) (which changes nothing) and follow(3, 4), user 4 posts tweet 8, and it appears first.

    Twitter t3 = new Twitter();
    System.out.println(t3.getNewsFeed(3));
    t3.follow(3, 3);
    t3.postTweet(3, 7);
    System.out.println(t3.getNewsFeed(3));
    t3.unfollow(3, 4);
    t3.follow(3, 4);
    System.out.println(t3.getNewsFeed(3));
    t3.postTweet(4, 8);
    System.out.println(t3.getNewsFeed(3));
    
    []
    [7]
    [7]
    [8, 7]
    

    Complexity

    postTweet, follow and unfollow: O(1). Each does a few hash-map lookups, one list append or one set operation.

    getNewsFeed: O(F log F), where F is the number of people the user follows. Building users takes O(F). The heap holds at most one entry per person, so at most F + 1 entries, and each push and pop costs O(log F). Filling the heap costs O(F log F) and the loop takes out at most 10 entries, so the number of tweets ever posted does not matter.

    Space: O(T + S) for T tweets and S follow relations stored, plus O(F) temporary space per feed request.

    Edge cases and common mistakes

    • Comparing tweet ids to find the newer tweet. Ids are arbitrary. Only the clock gives the real order.
    • Forgetting the user's own tweets. The feed includes them even when the user follows nobody.
    • Storing a self-follow. The user would then be in users twice and every one of their tweets would show twice. That is why follow returns early.
    • Pushing every tweet into the heap. That brings back the cost of looking at all tweets. Push one per person and refill from that person's list.
    • Asking for a user with no list. Check tweets.containsKey(user) before reading the list, because tweets.get(user) returns null for a user who never posted.
    • Unfollowing and old tweets. The followee's tweets stay in tweets. They are just not collected into this user's feed any more, and they return if the user follows again.
    Keep going with Java500+ step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free