51 lines
1.4 KiB
Python
51 lines
1.4 KiB
Python
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))
|