LC 0355 [M] Design Twitter - ALawliet/algorithms GitHub Wiki

class Twitter:

    def __init__(self):
        self.count = 0
        self.tweetMap = defaultdict(list) # userId -> list of [count, tweetId]
        self.followMap = defaultdict(set) # userId -> set of followeeId
        

    def postTweet(self, userId: int, tweetId: int) -> None:
        self.tweetMap[userId].append([self.count, tweetId])
        self.count -= 1 # reverse count for maxHeap
        

    def getNewsFeed(self, userId: int) -> List[int]:
        res = [] # ordered starting from recent
        maxHeap = [] 
        
        self.followMap[userId].add(userId) # including themself
        
        for followeeId in self.followMap[userId]: # followed
            if followeeId in self.tweetMap: # followed tweets
                index = len(self.tweetMap[followeeId]) - 1 # followed last tweet
                count, tweetId = self.tweetMap[followeeId][index]
                maxHeap.append([count, tweetId, followeeId, index - 1])
                
        heapq.heapify(maxHeap) # their followed tweets in order
        
        while maxHeap and len(res) < 10: # last 10 of all followed tweets
            count, tweetId, followeeId, index = heapq.heappop(maxHeap) # most recent occured
            res.append(tweetId)
            if index >= 0: # person had more tweets, we popped it off so add it back
                count, tweetId = self.tweetMap[followeeId][index]
                heapq.heappush(maxHeap, [count, tweetId, followeeId, index - 1])
        return res
        

    def follow(self, followerId: int, followeeId: int) -> None:
        self.followMap[followerId].add(followeeId)
        

    def unfollow(self, followerId: int, followeeId: int) -> None:
        # if followeeId in self.followMap[followerId]:
        self.followMap[followerId].discard(followeeId)
        


# Your Twitter object will be instantiated and called as such:
# obj = Twitter()
# obj.postTweet(userId,tweetId)
# param_2 = obj.getNewsFeed(userId)
# obj.follow(followerId,followeeId)
# obj.unfollow(followerId,followeeId)