Start learning free

LESSON 409 OF 574 · HEAP

Design Twitter in Swift: News Feed with a Max-Heap Merge

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

    SWIFT · HEAP

    Design Twitter in Swift, explained

    Problem overview

    The class Twitter supports four operations: postTweet(userId, tweetId), follow(followerId, followeeId), unfollow(followerId, followeeId) and getNewsFeed(userId). A user's news feed is the ten most recent tweets posted by the user and by everyone the user follows, newest first. A feed with fewer than ten tweets returns them all, and an empty feed returns an empty array.

    Following goes one way only. Following someone you already follow, following yourself and unfollowing someone you don't follow all change nothing, and your own tweets are always in your feed.

    Here is the first example. User 1 posts tweet 5, and the feed is [5]. User 1 follows user 2, who posts tweet 6, and the feed becomes [6, 5], because 6 is newer. After unfollow(1, 2) the feed is [5] again. The statement asks for O(1) time in postTweet, follow and unfollow, and O(F log F) in getNewsFeed, where F is the number of people the user follows.

    Approach: a clock, one list per user, and a merge

    The simple design stores every tweet in one big array. Then getNewsFeed filters it for the user's people and takes the last ten. That is easy, but each feed request walks through every tweet ever posted, so the cost grows with the whole site instead of with the user.

    A better design splits the data into three parts:

    • a global clock time that counts every tweet posted. Tweet ids say nothing about age. In the second example, tweet 106 is newer than tweet 201, so each tweet is stored with its time;
    • tweets, a dictionary from a user id to that user's tweets as (time, tweetId) pairs. Times only grow, so each list is already sorted from oldest to newest, and the newest tweet is at the end;
    • following, a dictionary from a user id to the Set of ids that user follows. A set ignores a repeated follow, and adds and removes in O(1).

    With that, postTweet, follow and unfollow are one dictionary update each.

    The feed is a merge of sorted lists. Each person's tweets are sorted by time, so the newest tweet of the feed is the newest of the people's last tweets. Put those last tweets, one per person, in a max-heap ordered by time. Take the top: it is the newest tweet in the feed. Then put that person's previous tweet into the heap, since it is now their newest tweet not shown yet. Repeat until there are ten tweets or the heap is empty.

    The heap never holds more than one entry per person, so there are at most F + 1 entries (the user plus the people followed). Building it and taking up to ten tweets out costs O(F log F) overall, whatever the number of tweets posted.

    Step by step

    1. Declare the clock time as a var set to 0.
    2. Declare tweets as an empty [Int: [(Int, Int)]]. A user gets a list only when they first post.
    3. Declare following as an empty [Int: Set<Int>].
    4. In postTweet, add 1 to time so every tweet gets its own, larger time.
    5. Append (time, tweetId) to tweets[userId, default: []]. The default: [] gives a new user an empty array to append to.
    6. In getNewsFeed, start users as an array holding userId, because your own tweets are always in your feed.
    7. Append the followees with append(contentsOf:) when following[userId] exists. follow never stores a user following themselves, so nobody appears twice.
    8. Create the heap newest as a MaxHeap<[Int]> with the comparison { $0[0] < $1[0] }. The lesson provides MaxHeap, which takes a closure that says when one element is smaller than another. Each entry is [time, user, index]: the tweet's time, its author, and its position in that author's array. It must be a var.
    9. Loop over users.
    10. Skip people who never posted with if let userTweets = tweets[user].
    11. Compute index, the last position: userTweets.count - 1.
    12. Insert the entry [userTweets[index].0, user, index]. Only one tweet per person goes in at first.
    13. Declare feed as an empty [Int].
    14. Loop while newest isn't empty and feed.count < 10.
    15. Extract the top entry into top.
    16. Read the user and the index from top[1] and top[2].
    17. Append the tweet's id, tweets[user]![index].1, to feed.
    18. If index > 0, insert that person's previous tweet with index - 1. It takes their place in the heap.
    19. Return feed.
    20. In follow, return at once if followerId == followeeId.
    21. Insert followeeId into following[followerId, default: []].
    22. In unfollow, remove followeeId from following[followerId] using ?., which does nothing if the user follows nobody.

    The solution in Swift

    class Twitter {
    var time = 0
    var tweets = [Int: [(Int, Int)]]()
    var following = [Int: Set<Int>]()

    func postTweet(_ userId: Int, _ tweetId: Int) {
        time += 1
        tweets[userId, default: []].append((time, tweetId))
    }

    func getNewsFeed(_ userId: Int) -> [Int] {
        var users = [userId]
        if let followees = following[userId] {
            users.append(contentsOf: followees)
        }
        var newest = MaxHeap<[Int]>({ $0[0] < $1[0] })
        for user in users {
            if let userTweets = tweets[user] {
                let index = userTweets.count - 1
                newest.insert([userTweets[index].0, user, index])
            }
        }
        var feed = [Int]()
        while !newest.isEmpty && feed.count < 10 {
            let top = newest.extract()
            let user = top[1], index = top[2]
            feed.append(tweets[user]![index].1)
            if index > 0 {
                newest.insert([tweets[user]![index - 1].0, user, index - 1])
            }
        }
        return feed
    }

    func follow(_ followerId: Int, _ followeeId: Int) {
        if followerId == followeeId {
            return
        }
        following[followerId, default: []].insert(followeeId)
    }

    func unfollow(_ followerId: Int, _ followeeId: Int) {
        following[followerId]?.remove(followeeId)
    }
}

    The Swift details worth a closer look:

    • default: in a subscript. tweets[userId, default: []].append(...) reads the array, or an empty one if the key is missing, appends to it, and stores the result back. A plain tweets[userId] returns an optional and could not be appended to directly.
    • Tuples as values. (Int, Int) is a tuple. .0 is the time and .1 is the tweet id. The heap holds plain [Int] entries and compares their first number, so no comparison of tuples is needed.
    • The heap's closure. { $0[0] < $1[0] } uses the shorthand argument names. Times are all different, so ties don't happen; the user and index only say where the tweet lives.
    • if let. if let followees = following[userId] runs its block only when the dictionary holds a value for the key, and names it. Dictionary lookups always return an optional.
    • ! on tweets[user]!. It unwraps the optional and crashes if there is no value. Here it is safe: the entry came from a tweet of that user, so the array exists.
    • Optional chaining. following[followerId]?.remove(followeeId) does nothing when there is no set. Set.remove returns the removed element or nil, and the line ignores it. Removing a value that isn't there is also harmless.

    Complexity

    Time. postTweet is one counter update and one append, O(1) amortized. follow and unfollow are one set insert or remove, O(1) on average. getNewsFeed inserts at most F + 1 entries and then does at most ten extracts and ten inserts, each on a heap of at most F + 1 entries, so it takes O(F log F).

    Space. O(T + E) for the stored data, where T is the number of tweets and E the number of follow relations, plus O(F) for one feed's heap.

    Edge cases and common mistakes

    • Using the tweet id to order tweets. Ids are arbitrary. With user 1 posting 101 to 106 and user 2 posting 201 to 206 in turns, the feed of user 1 is [206, 106, 205, 105, 204, 104, 203, 103, 202, 102], ordered by when the tweets were posted. The clock stores 101 at time 1 and 201 at time 2.
    • More than ten tweets. The same example has twelve tweets in user 1's feed, and only the ten newest are returned. User 2 doesn't follow user 1, so their feed is [206, 205, 204, 203, 202, 201].
    • An empty feed. If nobody in the feed has posted, getNewsFeed returns []. The loop over users skips people without tweets, and the while loop never runs.
    • Following yourself. Without the early return in follow, a user would be in users twice and every one of their tweets would appear twice. With it, follow(3, 3) followed by postTweet(3, 7) gives the feed [7].
    • Unfollowing someone you don't follow. It is a no-op. After unfollow(3, 4) and then follow(3, 4), user 3 follows user 4 even before user 4 has posted. Once user 4 posts tweet 8, the feed of user 3 is [8, 7].
    • Old tweets of a new followee. Following shows the followee's tweets from before, too. The lists are kept per user, so nothing needs copying when you follow.
    • Following twice. A set stores the id once. With users 9 and 8 followed by user 1 (9 twice), the feed contains each tweet once and following[1] holds two ids.
    • Collecting everything. Gathering all tweets of all followees and sorting them works for the examples, but takes time proportional to the number of tweets each time.

    Testing

    • postTweet(1, 5), feed of 1 gives [5]; follow(1, 2), postTweet(2, 6), feed of 1 gives [6, 5]; unfollow(1, 2), feed of 1 gives [5].
    • On a new Twitter, getNewsFeed(3) gives [].
    • On a new Twitter, postTweet(9, 40), postTweet(8, 30) and postTweet(9, 10), then user 1 follows 9, 8 and 9 again: the feed of user 1 is [10, 30, 40].
    Keep going with Swift500+ step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free