Start learning free

LESSON 408 OF 573 · HEAP

Design Twitter in Kotlin: News Feed by Merging 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 23
    Ctrl + Enter to check

    KOTLIN · HEAP

    Design Twitter in Kotlin, explained

    Problem overview

    You write a class Twitter with four methods. postTweet(userId, tweetId) records a tweet. follow(followerId, followeeId) and unfollow(followerId, followeeId) change who sees whose tweets. getNewsFeed(userId) returns the ids of the 10 most recent tweets written by the user or by anyone the user follows, newest first, as a List<Int>.

    The rules around the edges: following is one way; following someone twice, following yourself, or unfollowing someone you don't follow changes nothing; your own tweets are always in your feed; and an empty feed is an empty list. Tweet ids say nothing about age. A call posts later than all calls before it, so "newest" means "posted last", even if its id is smaller.

    Approach

    The simple solution stores every tweet in one list with its time. To build a feed, it scans the list, keeps the tweets whose author is in the user's group, and takes the last 10. Or it collects the group's tweets and sorts them. Either way the cost of getNewsFeed grows with every tweet ever posted, including the tweets of people nobody in the feed follows.

    The better design stores each user's tweets in their own list, oldest first. Because time only grows, a new tweet always goes at the end, so each list is already sorted. A feed is then a merge of several sorted lists, each read from its newest end, stopping after 10 items. This is the standard way to merge k sorted lists: a heap holds the next candidate of each list.

    Concretely, the heap starts with the newest tweet of each person in the feed. The top is the newest tweet overall, so it goes into the feed. Then that person's next older tweet takes its place in the heap. After 10 rounds, or when the heap is empty, the feed is complete. The heap never holds more than one entry per person, and at most 10 tweets are ever taken out.

    Step by step

    1. Three fields. time is a global clock (a var, starting at 0); tweets maps each user id to a list of IntArray pairs holding a tweet's time and id, oldest first; following maps each user id to the set of ids they follow. A set ignores repeated follows and adds and removes in O(1).
    2. Posting. postTweet increases time, creates the user's list if this is their first tweet, and appends intArrayOf(time, tweetId). Each user's list stays sorted by time without any sorting.
    3. Who is in the feed. getNewsFeed starts a list users with the user, since your own tweets are always shown, and adds everyone the user follows.
    4. Seed the heap. newest is a PriorityQueue<IntArray> built with compareByDescending { it[0] }, so the entry with the biggest time is on top. For each person who has posted, an entry intArrayOf(time, user, index) describes their newest tweet: its time, its author and its position in the author's list. People who never posted are skipped.
    5. Take up to 10. While the heap is not empty and the feed has fewer than 10 ids, poll the top entry, unpack it with val (_, user, index) = …, and add that tweet's id to feed.
    6. Refill from the same author. If index > 0, the author has an older tweet. Push it with index - 1: it is now their newest tweet not yet shown.
    7. Return the feed. It is already newest first.
    8. Follow and unfollow. follow returns at once when the two ids are equal, creates the follower's set when needed, and adds the followee. unfollow removes the followee when the follower has a set at all. Removing a missing value does nothing.

    The solution in Kotlin

    class Twitter {
        var time = 0
        val tweets = mutableMapOf<Int, MutableList<IntArray>>()
        val following = mutableMapOf<Int, MutableSet<Int>>()
    
        fun postTweet(userId: Int, tweetId: Int) {
            time++
            if (!tweets.containsKey(userId)) {
                tweets[userId] = mutableListOf()
            }
            tweets.getValue(userId).add(intArrayOf(time, tweetId))
        }
    
        fun getNewsFeed(userId: Int): List<Int> {
            val users = mutableListOf(userId)
            if (following.containsKey(userId)) {
                users.addAll(following.getValue(userId))
            }
            val newest = PriorityQueue<IntArray>(compareByDescending { it[0] })
            for (user in users) {
                if (tweets.containsKey(user)) {
                    val index = tweets.getValue(user).size - 1
                    newest.add(intArrayOf(tweets.getValue(user)[index][0], user, index))
                }
            }
            val feed = mutableListOf<Int>()
            while (newest.isNotEmpty() && feed.size < 10) {
                val (_, user, index) = newest.poll()
                feed.add(tweets.getValue(user)[index][1])
                if (index > 0) {
                    newest.add(intArrayOf(tweets.getValue(user)[index - 1][0], user, index - 1))
                }
            }
            return feed
        }
    
        fun follow(followerId: Int, followeeId: Int) {
            if (followerId == followeeId) {
                return
            }
            if (!following.containsKey(followerId)) {
                following[followerId] = mutableSetOf()
            }
            following.getValue(followerId).add(followeeId)
        }
    
        fun unfollow(followerId: Int, followeeId: Int) {
            if (following.containsKey(followerId)) {
                following.getValue(followerId).remove(followeeId)
            }
        }
    }
    

    Kotlin details worth knowing:

    • getValue versus []. For a map, map[key] gives a nullable value (null for a missing key). getValue(key) returns the plain value and throws if the key is missing. The code uses it only after a containsKey check.
    • Mutable collections. mutableMapOf, mutableListOf and mutableSetOf create collections that can change. A val field still always points at the same map, but the map's contents grow. The empty mutableListOf() and mutableSetOf() need no type argument, because the map's declared type tells Kotlin what they hold.
    • Destructuring an array. val (_, user, index) = newest.poll() takes the first three elements of the IntArray. The underscore skips the time, which has already done its job.
    • The comparator. compareByDescending { it[0] } compares two entries by their first number, largest first. Only times decide the order, and every tweet has a different time, so ties do not occur.
    • Method names. Kotlin's PriorityQueue is Java's, so poll() removes the top and returns it, and isNotEmpty() works on it like on any collection.

    Seeing it run

    This code uses the class the way the lesson's tests do: it creates a Twitter, posts, follows and prints feeds.

    fun main() {
        val twitter = Twitter()
        twitter.postTweet(1, 5)
        println(twitter.getNewsFeed(1))
        twitter.follow(1, 2)
        twitter.postTweet(2, 6)
        println(twitter.getNewsFeed(1))
        twitter.unfollow(1, 2)
        println(twitter.getNewsFeed(1))
    
        val ids = Twitter()
        ids.postTweet(1, 900)
        ids.postTweet(2, 5)
        ids.follow(3, 1)
        ids.follow(3, 2)
        println(ids.getNewsFeed(3))
    
        val busy = Twitter()
        busy.follow(1, 2)
        for (i in 1..6) {
            busy.postTweet(1, 100 + i)
            busy.postTweet(2, 200 + i)
        }
        println(busy.getNewsFeed(1))
        println(busy.getNewsFeed(2))
    
        val odd = Twitter()
        println(odd.getNewsFeed(3))
        odd.follow(3, 3)
        odd.postTweet(3, 7)
        println(odd.getNewsFeed(3))
        odd.unfollow(3, 4)
        odd.follow(3, 4)
        odd.postTweet(4, 8)
        println(odd.getNewsFeed(3))
    }
    

    It prints:

    [5]
    [6, 5]
    [5]
    [5, 900]
    [206, 106, 205, 105, 204, 104, 203, 103, 202, 102]
    [206, 205, 204, 203, 202, 201]
    []
    [7]
    [8, 7]
    

    The first three lines show a follow and an unfollow: tweet 6 appears first because it is newer, then disappears again. The fourth line shows that the clock decides the order: tweet 5 was posted after tweet 900, so it comes first. Lines five and six show the 10-tweet cut: user 1 sees 12 tweets but only the 10 newest, while user 2, who follows nobody, sees the 6 tweets they wrote. The last three lines cover the odd cases: an empty feed prints [], following yourself changes nothing, and following someone before they post is fine, because their later tweet shows up.

    Complexity

    • postTweet, follow, unfollow: O(1) on average. Each is a few hash lookups plus a list append or a set change.
    • getNewsFeed: with F the number of people in the feed, building the heap takes F insertions of O(log F) each, and at most 10 rounds of one removal and one insertion add 10 · O(log F). The total is O(F log F), whatever the number of tweets.
    • Space: O(T + L) for T tweets and L follow relations, plus O(F) for one feed call.

    Edge cases and common mistakes

    • Ordering by tweet id. Ids are arbitrary. Only the clock tells which tweet is newer.
    • Following yourself. Without the early return, the user would appear twice in users and their tweets would be listed twice.
    • Fewer than 10 tweets. The loop ends when the heap is empty, so the feed is shorter, or empty.
    • Merging everything and sorting. It gives correct answers but slows down with every tweet in the system, not just with the feed's size.
    • Pushing every tweet of every author into the heap. One entry per author is enough, because an author's older tweet cannot be shown before their newer one.
    • Unfollowing without a set. A user who never followed anyone has no set. The containsKey check avoids touching a missing one.
    Keep going with Kotlin500+ step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free