Design Twitter — HashMap + Heap for News Feed

Sanjeev SharmaSanjeev Sharma
2 min read

Advertisement

Problem 203 · Design Twitter

Difficulty: Hard · Pattern: HashMap + Heap

Design: postTweet(userId, tweetId), getNewsFeed(userId) (10 most recent across followees), follow/unfollow.

Solutions

# Python
import heapq
from collections import defaultdict

class Twitter:
    def __init__(self):
        self.time = 0
        self.tweets = defaultdict(list)  # userId -> [(time, tweetId)]
        self.following = defaultdict(set)

    def postTweet(self, userId, tweetId):
        self.tweets[userId].append((self.time, tweetId))
        self.time += 1

    def getNewsFeed(self, userId):
        # Merge last 10 tweets from self + followees
        heap = []  # max-heap by negated time
        users = self.following[userId] | {userId}
        for u in users:
            lst = self.tweets[u]
            if lst:
                idx = len(lst) - 1
                t, tid = lst[idx]
                heapq.heappush(heap, (-t, tid, u, idx))
        feed = []
        while heap and len(feed) < 10:
            t, tid, u, idx = heapq.heappop(heap)
            feed.append(tid)
            if idx > 0:
                nt, ntid = self.tweets[u][idx-1]
                heapq.heappush(heap, (-nt, ntid, u, idx-1))
        return feed

    def follow(self, followerId, followeeId):
        self.following[followerId].add(followeeId)

    def unfollow(self, followerId, followeeId):
        self.following[followerId].discard(followeeId)
// Java
class Twitter {
    int time = 0;
    Map<Integer, List<int[]>> tweets = new HashMap<>();
    Map<Integer, Set<Integer>> following = new HashMap<>();

    public void postTweet(int userId, int tweetId) {
        tweets.computeIfAbsent(userId, k->new ArrayList<>()).add(new int[]{time++, tweetId});
    }

    public List<Integer> getNewsFeed(int userId) {
        PriorityQueue<int[]> pq = new PriorityQueue<>((a,b)->b[0]-a[0]);
        Set<Integer> users = new HashSet<>(following.getOrDefault(userId, new HashSet<>()));
        users.add(userId);
        for (int u : users) {
            List<int[]> lst = tweets.getOrDefault(u, Collections.emptyList());
            if (!lst.isEmpty()) {
                int idx = lst.size()-1;
                pq.offer(new int[]{lst.get(idx)[0], lst.get(idx)[1], u, idx});
            }
        }
        List<Integer> feed = new ArrayList<>();
        while (!pq.isEmpty() && feed.size() < 10) {
            int[] cur = pq.poll();
            feed.add(cur[1]);
            if (cur[3] > 0) {
                int[] prev = tweets.get(cur[2]).get(cur[3]-1);
                pq.offer(new int[]{prev[0], prev[1], cur[2], cur[3]-1});
            }
        }
        return feed;
    }
    public void follow(int a, int b) { following.computeIfAbsent(a, k->new HashSet<>()).add(b); }
    public void unfollow(int a, int b) { following.getOrDefault(a, new HashSet<>()).remove(b); }
}

Complexity

  • postTweet: O(1)
  • getNewsFeed: O(f log f + 10 log f) where f = followee count
  • follow/unfollow: O(1)

Advertisement

Sanjeev Sharma

Written by

Sanjeev Sharma

Full Stack Engineer · E-mopro