import heapq class Twitter: def __init__(self): self.time = 0 self.tweets = {} self.follows = {} def postTweet(self, userId, tweetId): self.time += 1 if userId not in self.tweets: self.tweets[userId] = [] self.tweets[userId].append((-self.time, tweetId)) def getNewsFeed(self, userId): heap = [] if userId in self.tweets: heap.extend(self.tweets[userId][-10:]) if userId in self.follows: for v in self.follows[userId]: if v in self.tweets: heap.extend(self.tweets[v][-10:]) heapq.heapify(heap) res = [] while heap and len(res) < 10: res.append(heapq.heappop(heap)[1]) return res def follow(self, followerId, followeeId): if followerId == followeeId: return if followerId not in self.follows: self.follows[followerId] = set() self.follows[followerId].add(followeeId) def unfollow(self, followerId, followeeId): if followerId in self.follows and followeeId in self.follows[followerId]: self.follows[followerId].remove(followeeId) if __name__ == "__main__": tw = Twitter() tw.postTweet(1, 5) print(tw.getNewsFeed(1)) tw.follow(1, 2) tw.postTweet(2, 6) print(tw.getNewsFeed(1)) tw.unfollow(1, 2) print(tw.getNewsFeed(1))