DSA Designing a Mini Twitter

The problem: the mini Twitter spec

A mini Twitter stores tweets and builds a personal news feed for each user. Three operations define it. A post adds a tweet under its author, a follow records that one user follows another, and a feed returns up to ten tweet ids, newest first.

A feed mixes the reader's own tweets with the tweets of every followed user. Feed order is by posting time, newest first. A global time counter stamps each post, so no two tweets share a stamp.

script = [
    "post(1, 101)",
    "post(2, 102)",
    "post(1, 103)",
    "post(3, 104)",
    "follow(1, 2)",
    "follow(1, 3)",
    "feed(1) -> [104, 103, 102, 101]",
    "feed(2) -> [102]",
    "feed(3) -> [104]",
]
for line in script:
    print(line)

User 1 follows users 2 and 3, so the feed of user 1 merges three authors. Users 2 and 3 follow nobody, so each of their feeds holds only their own tweets.

Intuition: gather every tweet and sort the pile

The first design builds a feed by collecting every tweet of every relevant author into one pile and sorting the pile by time. The result is correct. The pile is the problem.

A busy author may hold thousands of tweets, and a feed needs only ten. Sorting the whole pile on every feed request pays for tweets the reader will never see.

tweets = {
    1: [(1, 101), (3, 103)],
    2: [(2, 102)],
    3: [(4, 104)],
}
pile = []
for author in sorted(tweets):
    pile.extend(tweets[author])
pile.sort(reverse=True)
print("tweets collected:", len(pile))
print("feed:", [tid for t, tid in pile[:10]])
print("tweets sorted but never shown:", len(pile) - min(10, len(pile)))

This toy pile is tiny, so nothing looks wasted. The waste scales with total tweets, not with the ten the feed shows. The real design reads only the heads of the author lists.

Brute force: the collect-and-sort feed, counted

The naive design is the intuition, written as a class. Each feed collects all tweets from the reader and the followed authors, sorts them newest first, and keeps ten. A counter records how many tweets each feed collects.

class NaiveTwitter:
    def __init__(self):
        self.time = 0
        self.tweets = {}
        self.follows = {}
        self.collected = 0

    def post(self, user, tweet):
        self.time += 1
        self.tweets.setdefault(user, []).append((self.time, tweet))

    def follow(self, a, b):
        self.follows.setdefault(a, set()).add(b)

    def feed(self, user):
        authors = set(self.follows.get(user, set()))
        authors.add(user)
        pile = []
        for author in sorted(authors):
            pile.extend(self.tweets.get(author, []))
        self.collected = len(pile)
        pile.sort(reverse=True)
        return [tid for t, tid in pile[:10]]

tw = NaiveTwitter()
tw.post(1, 101)
tw.post(2, 102)
tw.post(1, 103)
tw.post(3, 104)
tw.follow(1, 2)
tw.follow(1, 3)
print("feed of 1:", tw.feed(1), "collected:", tw.collected)
print("feed of 2:", tw.feed(2), "collected:", tw.collected)
print("feed of 3:", tw.feed(3), "collected:", tw.collected)
for i in range(11, 31):
    tw.post(2, 100 + i)
print("feed of 1 after 20 more posts by user 2:", tw.feed(1), "collected:", tw.collected)

The feeds match the script. The collected count tells the cost story. Twenty extra posts by one followed author raise every later feed of user 1 by twenty collected tweets, though the feed still shows ten.

The optimal solution: merge list heads with a heap

Each author's list already sits in time order with its newest tweet at the end, and a heap over the current heads of those lists hands back the newest remaining tweet on every pop. That merge is the whole design.

A heap is a structure that returns its smallest entry on each pop. Times enter negated, so the smallest entry is the newest tweet. After a pop, the next tweet from the same author enters the heap, and the merge stops after ten pops.

import heapq

class Twitter:
    def __init__(self):
        self.time = 0
        self.tweets = {}
        self.follows = {}

    def post(self, user, tweet):
        self.time += 1
        self.tweets.setdefault(user, []).append((self.time, tweet))

    def follow(self, a, b):
        self.follows.setdefault(a, set()).add(b)

    def feed(self, user):
        authors = set(self.follows.get(user, set()))
        authors.add(user)
        heap = []
        for author in sorted(authors):
            lst = self.tweets.get(author, [])
            if lst:
                idx = len(lst) - 1
                t, tid = lst[idx]
                heapq.heappush(heap, (-t, tid, author, idx))
        out = []
        while heap and len(out) < 10:
            negt, tid, author, idx = heapq.heappop(heap)
            out.append(tid)
            if idx > 0:
                t2, tid2 = self.tweets[author][idx - 1]
                heapq.heappush(heap, (-t2, tid2, author, idx - 1))
        return out

    def feed_of(self, user):
        return sorted(self.tweets.get(user, []), reverse=True)

tw = Twitter()
tw.post(1, 101)
tw.post(2, 102)
tw.post(1, 103)
tw.post(3, 104)
tw.follow(1, 2)
tw.follow(1, 3)
print("feed of 1:", tw.feed(1))
print("feed of 2:", tw.feed(2))
print("feed of 3:", tw.feed(3))

Every feed matches the script from the first section. The heap never holds more than one entry per author, so its size tracks the follow count, not the tweet count.

Walkthrough: building the feed of user 1

The walk reruns the script and prints the state after the posts and follows. It then pops the heap step by step, showing which author's next tweet enters after each pop.

import heapq

tweets = {1: [(1, 101), (3, 103)], 2: [(2, 102)], 3: [(4, 104)]}
for author in sorted(tweets):
    print("author", author, "tweets in time order:", tweets[author])
heap = []
for author in sorted(tweets):
    lst = tweets[author]
    idx = len(lst) - 1
    t, tid = lst[idx]
    heapq.heappush(heap, (-t, tid, author, idx))
print("heap heads as (time, tweet):", sorted((-negt, tid) for negt, tid, a, i in heap))
out = []
while heap and len(out) < 10:
    negt, tid, author, idx = heapq.heappop(heap)
    out.append(tid)
    print("pop tweet", tid, "from author", author)
    if idx > 0:
        t2, tid2 = tweets[author][idx - 1]
        heapq.heappush(heap, (-t2, tid2, author, idx - 1))
        print("  author", author, "offers next tweet", tid2)
print("feed:", out)

The pops come out in time order, 4, 3, 2, 1. Each author appears in the heap at most once at any moment, because a refill enters only after that author's head leaves.

Complexity

A post appends in O(1) time and a follow adds to a set in O(1) time. A feed makes at most ten pops and ten pushes on a heap of one entry per author, so its time is O(log a) per shown tweet for a authors, and its extra space is O(a).

for posts_per_author in [10, 100, 1000]:
    authors = 3
    collected = authors * posts_per_author
    heap_pops = 10
    print("posts per author", posts_per_author, "- naive collects:", collected, " heap pops per feed:", heap_pops)

History growth hits only the naive column. The heap column stays at ten pops because the feed stops at ten tweets, however long the lists behind it grow.

Edge cases and pitfalls

A user with no tweets and no follows gets an empty feed. Following the same user twice changes nothing, because the follow graph is a set. A post made after a follow still appears, since the feed reads the lists live at request time.

The classic pitfalls are order and determinism. Printing a follow set in raw order would vary between runs, so any printed set is sorted first. Feed order itself is fixed by the unique time stamps, never by set order.

import heapq

class Twitter:
    def __init__(self):
        self.time = 0
        self.tweets = {}
        self.follows = {}

    def post(self, user, tweet):
        self.time += 1
        self.tweets.setdefault(user, []).append((self.time, tweet))

    def follow(self, a, b):
        self.follows.setdefault(a, set()).add(b)

    def feed(self, user):
        authors = set(self.follows.get(user, set()))
        authors.add(user)
        heap = []
        for author in sorted(authors):
            lst = self.tweets.get(author, [])
            if lst:
                idx = len(lst) - 1
                t, tid = lst[idx]
                heapq.heappush(heap, (-t, tid, author, idx))
        out = []
        while heap and len(out) < 10:
            negt, tid, author, idx = heapq.heappop(heap)
            out.append(tid)
            if idx > 0:
                t2, tid2 = self.tweets[author][idx - 1]
                heapq.heappush(heap, (-t2, tid2, author, idx - 1))
        return out

    def feed_of(self, user):
        return sorted(self.tweets.get(user, []), reverse=True)

tw = Twitter()
print("empty feed:", tw.feed(9))
tw.post(1, 101)
tw.follow(1, 1)
tw.follow(1, 2)
tw.follow(1, 2)
print("follows of 1, sorted:", sorted(tw.follows[1]))
print("feed of 1:", tw.feed(1))
tw.post(2, 102)
print("feed of 1 after a later post by 2:", tw.feed(1))
for i in range(12):
    tw.post(1, 200 + i)
feed = tw.feed(1)
print("feed length with 14 visible tweets:", len(feed))
print("feed:", feed)

The long feed stops at exactly ten ids. The cap is part of the spec, and the heap enforces it by stopping the pops.