Design Twitter

MediumHeapLeetCode 355 ↗World 9-11
0:00 / 0:00

Build a tiny Twitter with post, follow and a news feed. Hash maps hold tweets and follows, and a max-heap merges the ten newest.

▼

The problem

LeetCode 355 (Medium). Design a small social feed with postTweet(userId, tweetId), follow(followerId, followeeId), unfollow(followerId, followeeId) and getNewsFeed(userId), which returns the ids of the 10 most recent tweets posted by the user or by anyone they follow, newest first. (The visuals call it a small social feed; the brand name appears only in the problem title.)

Example (LeetCode's): postTweet(1, 5); getNewsFeed(1) → [5]; follow(1, 2); postTweet(2, 6); getNewsFeed(1) → [6, 5]; unfollow(1, 2); getNewsFeed(1) → [5].

TRY IT ON LEETCODE ▶

The solution

import heapq
from collections import defaultdict

class Twitter:
    def __init__(self):
        self.time = 0
        self.tweets = defaultdict(list)  # (time, id)
        self.follows = defaultdict(set)
    def postTweet(self, user, tweet):
        self.time += 1
        self.tweets[user].append((self.time, tweet))
    def follow(self, a, b): self.follows[a].add(b)
    def unfollow(self, a, b): self.follows[a].discard(b)
    def getNewsFeed(self, user):
        heap = []
        for u in self.follows[user] | {user}:
            if self.tweets[u]:
                i = len(self.tweets[u]) - 1
                time, tid = self.tweets[u][i]
                heap.append((-time, tid, u, i))
        heapq.heapify(heap)
        feed = []
        while heap and len(feed) < 10:
            _, tid, u, i = heapq.heappop(heap)
            feed.append(tid)
            if i > 0:
                time, tid = self.tweets[u][i - 1]
                heapq.heappush(heap, (-time, tid, u, i - 1))
        return feed

Transcript

Design Twitter. Build a small social feed. Users post tweets, follow and unfollow each other, and fetch a news feed: the ten most recent tweet ids from themselves and everyone they follow, newest first.

User one posts tweet five, and their feed is five. User one follows user two, who posts tweet six. Now the feed is six, five. After the unfollow, it's just five.

The naive feed gathers every followee's tweets and sorts them all. That's order T log T for T tweets, just to keep ten.

Better: stamp each post with a global clock. One hash map sends each user to their tweets, as time and id pairs in posting order. Another sends them to the set of users they follow. Post appends, follow adds, unfollow removes, each order one.

Each list is already sorted, so merge them. Seed a max-heap with every followee's newest tweet, yourself included. Pop the newest into the feed, then push that user's next older tweet. Stop after ten pops.

In code: post ticks the clock and appends. The feed seeds the heap, then pops up to ten times, pushing each user's previous tweet.

On the example: tweet five gets time one, so the feed is five. User one follows two, and tweet six gets time two. The heap holds six and five: pop six, then five. After the unfollow, only five.

Post, follow and unfollow are order one. A feed is order F to seed the heap, plus ten pops of log F. Space is order of all tweets and follows.

Stamp every post, keep a set of follows, and merge with a heap. That's Design Twitter.