MediumLeetCode 355Design — hash maps + heap merge

Design Twitter (LeetCode 355)

Problem 6 of 10 in Design a Data Structure — Practice Problems

The problem

Design a simplified Twitter with four operations:

  • postTweet(userId, tweetId): the user posts a tweet;
  • getNewsFeed(userId): return up to 10 most recent tweet IDs from the user and everyone they follow, newest first;
  • follow(followerId, followeeId) and unfollow(followerId, followeeId).
postTweet(1, 5)
getNewsFeed(1) -> [5]
follow(1, 2)
postTweet(2, 6)
getNewsFeed(1) -> [6, 5]   (6 is newer)
unfollow(1, 2)
getNewsFeed(1) -> [5]

How to think about it

Start with brute force. Keep one global list of every tweet as (user, tweetId) in posting order. For a feed, walk it backwards and keep tweets whose author is the user or a followee, stopping at 10. That is simple, but a feed can scan the entire history of the site to find 10 relevant tweets: O(total tweets) per call.

Split the data by user. Store each user’s tweets in their own list, and store each user’s followees in a set. Now a feed only needs to look at the lists of the user and the people they follow. Following and unfollowing become O(1) set operations.

The key insight: a feed is a k-way merge of recent tweets. Each user’s list is already in time order. The feed wants the 10 newest across several such lists, so you only ever need the last 10 tweets of each relevant user: nothing older can make the top 10. Tag every tweet with a global, increasing timestamp so tweets from different users can be compared, then use a heap (heapq.nlargest(10, ...)) to pick the 10 largest timestamps from those candidates.

Which pattern is this? Design a data structure: hash maps for lookup + a heap for ordering. The tell is “top k most recent across many sources”, which is the same shape as merging k sorted lists.

Approach

  1. Keep a counter time, a map tweets[user] → list of (time, tweetId), and a map follows[user] → set of followees.
  2. postTweet: increase time and append (time, tweetId) to the user’s list.
  3. getNewsFeed: take the user’s followees plus the user. Collect the last 10 tweets from each of their lists. Return the tweet IDs of the 10 largest by time, newest first.
  4. follow adds to the set; unfollow uses discard, so unfollowing someone you don’t follow is harmless.

Solution

from collections import defaultdict
import heapq

class Twitter:
    def __init__(self):
        self.time = 0; self.tweets = defaultdict(list); self.follows = defaultdict(set)
    def postTweet(self, uid, tid):
        self.time += 1; self.tweets[uid].append((self.time, tid))    # global timestamp
    def getNewsFeed(self, uid):
        users = self.follows[uid] | {uid}; cand = []                  # always include yourself
        for u in users: cand += self.tweets[u][-10:]                 # only the 10 newest can matter
        return [tid for _, tid in heapq.nlargest(10, cand)]
    def follow(self, a, b): self.follows[a].add(b)
    def unfollow(self, a, b): self.follows[a].discard(b)
# postTweet(1,5) feed(1)->[5] follow(1,2) postTweet(2,6) feed(1)->[6,5] unfollow(1,2) feed(1)->[5]

Dry run

The example sequence:

Operation Returns tweets follows[1]
postTweet(1, 5) — {1: [(1, 5)]} {}
getNewsFeed(1) [5] {1: [(1, 5)]} {}
follow(1, 2) — {1: [(1, 5)]} {2}
postTweet(2, 6) — {1: [(1, 5)], 2: [(2, 6)]} {2}
getNewsFeed(1) [6, 5] same {2}
unfollow(1, 2) — same {}
getNewsFeed(1) [5] same {}

In the second feed, the candidates are (1, 5) and (2, 6). Timestamp 2 is larger, so tweet 6 comes first.

Complexity

  • postTweet, follow, unfollow: O(1) (amortised append and set operations).
  • getNewsFeed: O(F) where F is the number of followees. It collects at most 10(F + 1) candidates, and heapq.nlargest(10, …) over c candidates is O(c log 10), which is O(F).
  • Space: O(T + E) for T tweets and E follow relationships.

Edge cases and pitfalls

  • Include the user’s own tweets even though they don’t “follow” themselves. The | {uid} handles it, and also stops a self-follow from producing duplicates.
  • Unfollowing someone you don’t follow must not crash: discard, not remove.
  • Fewer than 10 tweets in total: return what exists.
  • Comparing by tweet ID instead of a timestamp is wrong: IDs are not guaranteed to increase with time.
  • Scaling follow-up: for users with millions of followers, real systems precompute feeds on write (“fan-out on write”) for most users and merge on read for celebrities. A true k-way merge with a heap of one pointer per list avoids copying 10 tweets from every followee.