design_twitter.py

import heapq


class Tweet:
    """
        Every tweet should have 
        1. an id to uniquely identify the tweet,
        2. the time tweet was posted
        3. a link to other tweets. 
    """   
    def __init__(self, tweet_id: int, time_stamp: int):
        self.tweet_id: int = tweet_id
        self.time_posted: int = time_stamp

        # No tweet has been added yet.
        self.next: Tweet = None


class PriorityTweetQueue:
    def __init__(self, capacity: int) -> None:
        self.capacity = capacity
        self.recent_tweets = []

    def add(self, tweet: Tweet) -> None:
        if len(self.recent_tweets) < self.capacity:
            # heapq will automatically sort in ascending order (i.e. the smallest
            # time_posted will be at the top of the queue) but since we need the 
            # largest time_posted to be at the top of the queue, we then negate 
            # the time_posted so the largest time_posted will be smallest
            # thereby will be at the top of the queue. 
            heapq.heappush(self.recent_tweets, (-tweet.time_posted, tweet))

    def pop(self):
        if not self.empty():
            time_posted, tweet = heapq.heappop(self.recent_tweets)
            return tweet
        return None

    def empty(self) -> bool:
        return len(self.recent_tweets) == 0


class User:
    """
    A User (or a tweeter) must have possess the following:
        1. Id to uniquely identify the user
        2. People that he follows (so that he would be able to view their
            recent tweets and also his.). To be able to view his/her own tweets,
            then he must follow himself/herself.
        3. Must be able to follow and unfollow other Users.
        4. Must be able to post tweets.
    """
    
    def __init__(self, user_id: int) -> None:
        self.user_id: int = user_id
        
        # Store users that user_id follows.
        self.follows = set()
        
        self.user_tweets = None
        
        # User follows himself/herself.
        self.follow(user_id)
    
    def follow(self, other_user_id: int) -> None:
        self.follows.add(other_user_id)
    
    def unfollow(self, other_user_id: int) -> None:
        # To unfollow, other_user_id must exist and I can't unfollow myself.
        if other_user_id in self.follows and other_user_id != self.user_id:
            self.follows.remove(other_user_id)
    
    def post(self, tweet_id: int, time_stamp: int) -> None:
        """
            Make sure the most recent tweet is at the head.
            => Nothing added yet.
                null
            => post(1)
                1 -> null
            => post(2)
                2 -> 1 -> null
        """
        new_tweet = Tweet(tweet_id, time_stamp)
        new_tweet.next = self.user_tweets
        self.user_tweets = new_tweet

        print(self)

    def __str__(self):
        tweets = ""
        user_tweets = self.user_tweets
        while user_tweets:
            tweets += f"({user_tweets.tweet_id}, {user_tweets.time_posted})"
            user_tweets = user_tweets.next
            if user_tweets:
                tweets += " -> "
        return f"{self.user_id}: {tweets}"


class Twitter:
    # Global Twitter time to differentiate all tweets.
    TIMESTAMP = 1

    def __init__(self):
        # Map of user_id and User
        self.users = {}
        self.recent_tweet_capacity = 10
        
    def postTweet(self, user_id: int, tweet_id: int) -> None:
        """
            For a user to post a tweet, then the user must exist
            hence create the user if the user doesn't exists,
            then user can post tweet.
        """
        if user_id not in self.users:
            new_user = User(user_id)
            self.users[user_id] = new_user

        self.users.get(user_id).post(tweet_id, Twitter.TIMESTAMP)
        Twitter.TIMESTAMP += 1
        
    def getNewsFeed(self, user_id: int) -> List[int]:
        if user_id not in self.users:
            return []

        """
            Now, if a user follows other users, then the user's news feed as well
            as the other users' news feed would be displayed on the user's home page.
            
            Hence, we need to 
            1. Get all the users that self current user follows 
               (we all know that self would also include the current user itself
                since a user follows itself).
            2. Insert all the user's tweet that the current user follows in a 
               priority queue based on each user's current tweet (which would be at 
               the head their respective tweets) time_stamp difference (tweet comparator) 
               as self would make sure the most recent tweet is at the head of the queue.
        """
        #  We technically can't have more than number of users that have posted a tweet.
        num_of_users: int = len(self.users)
        
        recent_tweets = PriorityTweetQueue(num_of_users)
        
        curr_user = self.users.get(user_id)
        # Store all the tweets of the users that the current user is following.
        for other_user_id in curr_user.follows:
            other_user_tweets = self.users.get(other_user_id).user_tweets
            if other_user_tweets is not None:
                recent_tweets.add(other_user_tweets)
        
        top_required_tweets = []
        pushed_tweet = 0
        
        while not recent_tweets.empty() and pushed_tweet < self.recent_tweet_capacity:
            top_tweet = recent_tweets.pop()

            top_required_tweets.append(top_tweet.tweet_id)
            
            next_tweet = top_tweet.next
            if next_tweet is not None:
                recent_tweets.add(next_tweet)
            
            pushed_tweet += 1
        
        return top_required_tweets

    def follow(self, follower_id: int, followee_id: int) -> None:
        # For followerId to follow followeeId, then both of them
        # must first exist.

        # 1. Add the follower.
        if follower_id not in self.users:
            follower = User(follower_id)
            self.users[follower_id] = follower
    
        # 2. Add the followee.
        if followee_id not in self.users:
            followee = User(followee_id)
            self.users[followee_id] = followee

        # Make the follower follow the followee.
        self.users.get(follower_id).follow(followee_id)
        
    def unfollow(self, follower_id: int, followee_id: int) -> None:
        # A follower can only unfollow a followee if the follower exists.
        if follower_id in self.users:
            # Note: user not unfollowing himself/herself has been
            # implemented in the User class.
            self.users.get(follower_id).unfollow(followee_id)

"""
Design a simplified version of Twitter where users can post tweets, follow/unfollow another user, 
and is able to see the 10 most recent tweets in the user's news feed.

Implement the Twitter class:

Twitter() Initializes your twitter object.
- postTweet(self, user_id: int, tweet_id: int): 
  Composes a new tweet with ID tweetId by the user userId. Each call to this function will be made with a unique tweetId.
- getNewsFeed(self, user_id: int) 
  Retrieves the 10 most recent tweet IDs in the user's news feed. 
  Each item in the news feed must be posted by users who the user followed or by the user themself. 
  Tweets must be ordered from most recent to least recent.
- follow(self, follower_id: int, followee_id: int): 
  The user with ID followerId started following the user with ID followeeId.
- unfollow(self, follower_id: int, followee_id: int):
  The user with ID followerId started unfollowing the user with ID followeeId.
 

Example 1:

Input
["Twitter", "postTweet", "getNewsFeed", "follow", "postTweet", "getNewsFeed", "unfollow", "getNewsFeed"]
[[], [1, 5], [1], [1, 2], [2, 6], [1], [1, 2], [1]]
Output
[null, null, [5], null, null, [6, 5], null, [5]]

Explanation
Twitter twitter = new Twitter();
twitter.postTweet(1, 5); // User 1 posts a new tweet (id = 5).
twitter.getNewsFeed(1);  // User 1's news feed should return a list with 1 tweet id -> [5]. return [5]
twitter.follow(1, 2);    // User 1 follows user 2.
twitter.postTweet(2, 6); // User 2 posts a new tweet (id = 6).
twitter.getNewsFeed(1);  // User 1's news feed should return a list with 2 tweet ids -> [6, 5]. Tweet id 6 should precede tweet id 5 because it is posted after tweet id 5.
twitter.unfollow(1, 2);  // User 1 unfollows user 2.
twitter.getNewsFeed(1);  // User 1's news feed should return a list with 1 tweet id -> [5], since user 1 is no longer following user 2.
 

Constraints:

1 <= userId, followerId, followeeId <= 500
0 <= tweetId <= 104
All the tweets have unique IDs.
At most 3 * 104 calls will be made to postTweet, getNewsFeed, follow, and unfollow.

"""