Problem statement

https://leetcode.com/problems/design-twitter/

Solution

For follow and unfollow we just keep defaultdicts with lists or sets, so we can easily iterate over all persons given user follow. For postTweet we need to keep list of all tweets this person made with timestamps. These 3 operations can be done in O(1).

The main difficulty with getNewsFeed(userID), we need to go over all K persons which follow userID and choose 10 of them with biggest timestamps. It can be done in O(10K) complexity, how can we do it? We kinda need K iterators, choose the smallest of them and repeat the procedure. We can just merge all these lists, create heap in O(10K) and then return 10 biggest elements in O(log K). However, complexity is still O(10K). (If we do not use heaps it will be O(100K)). Open question, can we do better?

Complexity

It is O(1) for first 3 functions and O(10K) for the last one, where K is numbers of person in feed.

Code

class Twitter:
    def init(self):
        self.foll = defaultdict(set)
        self.tweets = defaultdict(list)
        self.time = 0
        
    def postTweet(self, userId, tweetId):
        self.tweets[userId].append((self.time, tweetId))
        self.time -= 1
        
    def getNewsFeed(self, userId):
        feed = self.foll[userId] | set([userId])
        tweets = list(chain(*[self.tweets[idx] for idx in feed]))
        heapify(tweets)
        ans = []
        for i in range(min(10, len(tweets))):
            ans.append(heappop(tweets)[1])
        return ans
        
    def follow(self, followerId, followeeId):
        self.foll[followerId].add(followeeId)
        
    def unfollow(self, followerId, followeeId):
        self.foll[followerId].discard(followeeId)