Design Twitter — HashMap + Heap for News Feed
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