Start learning free

LESSON 411 OF 576 · HEAP

Design Twitter in Rust: A Max-Heap Merge of Each Feed

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 22
    Ctrl + Enter to check

    RUST · HEAP

    Design Twitter in Rust, explained

    Problem overview

    The struct Twitter supports five operations: Twitter::new() creates an empty service, post_tweet(user_id, tweet_id) posts a tweet, follow(follower_id, followee_id) and unfollow(follower_id, followee_id) change who sees whose tweets, and get_news_feed(user_id) returns a Vec<i32> with the ids of the 10 newest tweets by the user and the people the user follows, newest first.

    Following goes one way only. Following someone twice, following yourself, or unfollowing someone you don't follow changes nothing, and your own tweets are always in your feed. A feed with fewer than 10 tweets returns them all, and an empty feed returns an empty vector.

    For example, user 1 posts tweet 5, follows user 2, and user 2 posts tweet 6. User 1's feed is then [6, 5], because 6 is newer. After unfollow(1, 2) the feed is [5] again.

    The statement asks for O(1) time for post_tweet, follow and unfollow, and O(F log F) for get_news_feed, where F is the number of people the user follows.

    Approach: merge each person's tweets with a heap

    The simple approach gathers every tweet of the user and of everyone they follow, sorts the lot by time and takes the first 10. That cost grows with every tweet ever posted, even though the feed shows only 10.

    Two facts make it cheaper. First, each person's tweets are already in order, because they are appended as they are posted. Second, the feed needs only the 10 newest, so it never has to look far into anyone's history. This is the same merge as combining sorted lists: the heap holds the next candidate of each list, and the largest candidate is the next tweet of the feed.

    Tweet ids don't tell which tweet is newer: in the second example, tweet 106 is newer than tweet 201, and ids can be arbitrary anyway. So the struct keeps a global clock. Every tweet gets the next time, so a bigger time always means a newer tweet.

    The structure has three parts: the clock, a map from each user to a vector of (time, tweet_id) pairs (oldest first), and a map from each user to the set of ids they follow. A set makes a repeated follow harmless and makes both follow and unfollow O(1).

    To build a feed, put the newest tweet of each person into a max-heap, as a tuple (time, user, index) where index is the tweet's position in that user's vector. Tuples compare field by field, and times are all different, so the top is always the newest tweet not yet shown. Pop it, add its id to the feed, and push the same person's previous tweet, at index - 1, if there is one. Stop at 10 ids or when the heap is empty.

    Take the second example: users 1 and 2 take turns posting 101 to 106 and 201 to 206, and user 1 follows user 2. The vectors hold times 1, 3, 5, 7, 9, 11 for user 1 and 2, 4, 6, 8, 10, 12 for user 2. The heap starts with each user's newest tweet, and the first three rounds go like this:

    start: [(11, 1, 5), (12, 2, 5)]
take tweet 206
heap: [(10, 2, 4), (11, 1, 5)]
take tweet 106
heap: [(9, 1, 4), (10, 2, 4)]
take tweet 205
heap: [(8, 2, 3), (9, 1, 4)]

    The heap lists are printed in ascending order, so the entry on top is the last one.

    Step by step

    1. Declare the clock time as an i32.
    2. Declare tweets, a HashMap<i32, Vec<(i32, i32)>> from user id to that user's (time, tweet_id) pairs.
    3. Declare following, a HashMap<i32, HashSet<i32>> from user id to the ids the user follows.
    4. In new, return the struct with the clock at 0 and two empty maps.
    5. In post_tweet, advance the clock by 1.
    6. Push the pair (self.time, tweet_id) onto the user's vector, creating the vector on the user's first post with entry(...).or_insert(Vec::new()). Times only grow, so each vector stays ordered from oldest to newest.
    7. In get_news_feed, start a vector of users with user_id itself.
    8. Extend it with the followed ids, if the user follows anyone. No one appears twice, because follow never stores a user following themselves.
    9. Create the max-heap newest.
    10. Loop over the users.
    11. Skip the users who never posted, by testing self.tweets.get(&user) with if let.
    12. Find the newest tweet's index: the vector's length minus 1.
    13. Push (time, user, index) into the heap, with the time read from the pair's field .0.
    14. Create the vector feed.
    15. Loop while the heap has entries and the feed has fewer than 10 ids.
    16. Pop the top tuple and unpack it, ignoring the time.
    17. Push that tweet's id, the pair's field .1, onto the feed.
    18. Push the same user's previous tweet if index > 0, so the heap again holds one candidate per person.
    19. Return feed.
    20. In follow, return early when a user would follow themselves.
    21. Insert the followee into the follower's set, creating the set on first use.
    22. In unfollow, remove the followee from the follower's set if the follower has one. Removing an id that isn't there does nothing.

    The solution in Rust

    use std::collections::{BinaryHeap, HashMap, HashSet};

struct Twitter {
    time: i32,
    tweets: HashMap<i32, Vec<(i32, i32)>>,
    following: HashMap<i32, HashSet<i32>>,
}

impl Twitter {
    fn new() -> Twitter {
        Twitter { time: 0, tweets: HashMap::new(), following: HashMap::new() }
    }

    fn post_tweet(&mut self, user_id: i32, tweet_id: i32) {
        self.time += 1;
        self.tweets.entry(user_id).or_insert(Vec::new()).push((self.time, tweet_id));
    }

    fn get_news_feed(&self, user_id: i32) -> Vec<i32> {
        let mut users = vec![user_id];
        if let Some(followees) = self.following.get(&user_id) {
            users.extend(followees);
        }
        let mut newest = BinaryHeap::new();
        for user in users {
            if let Some(user_tweets) = self.tweets.get(&user) {
                let index = user_tweets.len() - 1;
                newest.push((user_tweets[index].0, user, index));
            }
        }
        let mut feed = Vec::new();
        while !newest.is_empty() && feed.len() < 10 {
            let (_, user, index) = newest.pop().unwrap();
            feed.push(self.tweets[&user][index].1);
            if index > 0 {
                newest.push((self.tweets[&user][index - 1].0, user, index - 1));
            }
        }
        feed
    }

    fn follow(&mut self, follower_id: i32, followee_id: i32) {
        if follower_id == followee_id {
            return;
        }
        self.following.entry(follower_id).or_insert(HashSet::new()).insert(followee_id);
    }

    fn unfollow(&mut self, follower_id: i32, followee_id: i32) {
        if let Some(followees) = self.following.get_mut(&follower_id) {
            followees.remove(&followee_id);
        }
    }
}

    Using it on the lesson's examples gives:

    post_tweet(1, 5)       feed(1) = [5]
follow(1, 2)
post_tweet(2, 6)       feed(1) = [6, 5]
unfollow(1, 2)         feed(1) = [5]

    With users 1 and 2 posting in turns as above, user 1's feed is 206 106 205 105 204 104 203 103 202 102 and user 2's feed, which includes only user 2's own tweets, is 206 205 204 203 202 201.

    The Rust details worth a closer look:

    • entry(...).or_insert(...). It returns a mutable reference to the value stored for the key, inserting the given value first when the key is missing. That is why .push(...) and .insert(...) can be chained onto it.
    • &self versus &mut self. get_news_feed only reads, so it takes &self, and the heap it builds is a local variable. post_tweet, follow and unfollow change the maps and take &mut self.
    • get and get_mut return an Option. if let Some(followees) = self.following.get(&user_id) runs its block only when the key is there. get_mut lends the set mutably, which remove needs.
    • users.extend(followees). followees is a &HashSet<i32>, which iterates over &i32. A vector of i32 can be extended from references to i32, because i32 is Copy.
    • Tuples in a BinaryHeap. BinaryHeap is a max-heap, and tuples compare field by field, so the tuple with the largest time is on top. The heap's element type is inferred from the first push.
    • Indexing a map. self.tweets[&user] panics when the key is missing. It is safe here, because the heap only holds users who have a vector.
    • The subtraction len() - 1. It is a usize and would wrap or panic for an empty vector. A vector exists only after a first post, so it always has at least one element.

    Complexity

    Time. post_tweet advances the clock and appends to a vector: O(1) on average. follow and unfollow are one set operation each, also O(1) on average. get_news_feed pushes at most F + 1 starting entries into the heap, which costs O(F log F), and then does at most 10 rounds of one pop and one push, each O(log F). The total is O(F log F), and it does not depend on how many tweets exist.

    Space. The maps hold every tweet and every follow relation, so memory is O(T + R) for T tweets and R follow relations. A single feed call needs O(F) extra space for the user list and the heap, plus the 10 ids it returns.

    Edge cases and common mistakes

    • Using tweet ids as times. After posting 900, 3 and 500 as user 1, the feed is [500, 3, 900]: newest first by posting order, not by id.
    • An empty feed. A user who has not posted and follows nobody gets an empty vector.
    • Following yourself. Without the early return, the user's id would be in the list twice and each of their tweets would show up twice. After follow(3, 3) and post_tweet(3, 7), the feed is just [7].
    • Following someone before they post. It works, because the feed reads tweets at the time of the call. After follow(3, 4) and a later post of tweet 8 by user 4, user 3's feed is [8, 7].
    • Repeated follows and useless unfollows. A second follow(3, 4) leaves the feed at [8, 7], and unfollow(3, 4) before any follow does nothing.
    • Pushing every tweet into the heap. It works but costs more with every tweet ever posted. Pushing one candidate per person is what bounds the cost.
    • Forgetting to push the previous tweet. Without it, each person would contribute at most one tweet to the feed.
    • More than 10 tweets. A user with 30 tweets (ids 0 to 29) gets 29 28 27 26 25 24 23 22 21 20.

    Testing

    • The first example gives the feeds 5, 6 5 and 5.
    • The second example gives 206 106 205 105 204 104 203 103 202 102 for user 1 and 206 205 204 203 202 201 for user 2.
    • The third example gives an empty line, then 7, 7 and 8 7.
    Keep going with Rust500+ step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free