355. Design Twitter

• Time: $O(n + k\log k)$, where $n = |\texttt{tweets}|$ and $k = \min(10, |\texttt{tweets}|))$
• Space: $O(n)$

C++JavaPython

struct Tweet {
  int id;
  int time;
  Tweet* next = nullptr;
};

struct User {
  int id;
  unordered_set<int> followeeIds;
  Tweet* tweetHead = nullptr;

  User() {}

  User(int id) : id(id) {
    follow(id);
  }

  void follow(int followeeId) {
    followeeIds.insert(followeeId);
  }

  void unfollow(int followeeId) {
    followeeIds.erase(followeeId);
  }

  void post(int tweetId, int time) {
    Tweet* oldTweetHead = tweetHead;
    tweetHead = new Tweet(tweetId, time);
    tweetHead->next = oldTweetHead;
  }
};

class Twitter {
 public:
  /** Compose a new tweet. */
  void postTweet(int userId, int tweetId) {
    if (!users.contains(userId))
      users[userId] = User(userId);
    users[userId].post(tweetId, time++);
  }

  /**
   * Retrieve 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
   * herself. Tweets must be ordered from most recent to least recent.
   */
  vector<int> getNewsFeed(int userId) {
    if (!users.contains(userId))
      return {};

    vector<int> newsFeed;

    auto compare = [](const Tweet* a, const Tweet* b) {
      return a->time < b->time;
    };
    priority_queue<Tweet*, vector<Tweet*>, decltype(compare)> maxHeap(compare);

    for (const int followeeId : users[userId].followeeIds) {
      Tweet* tweetHead = users[followeeId].tweetHead;
      if (tweetHead != nullptr)
        maxHeap.push(tweetHead);
    }

    int count = 0;
    while (!maxHeap.empty() && count++ < 10) {
      Tweet* tweet = maxHeap.top();
      maxHeap.pop();
      newsFeed.push_back(tweet->id);
      if (tweet->next)
        maxHeap.push(tweet->next);
    }

    return newsFeed;
  }

  /**
   * Follower follows a followee.
   * If the operation is invalid, it should be a no-op.
   */
  void follow(int followerId, int followeeId) {
    if (followerId == followeeId)
      return;
    if (!users.contains(followerId))
      users[followerId] = User(followerId);
    if (!users.contains(followeeId))
      users[followeeId] = User(followeeId);
    users[followerId].follow(followeeId);
  }

  /**
   * Follower unfollows a followee.
   * If the operation is invalid, it should be a no-op.
   */
  void unfollow(int followerId, int followeeId) {
    if (followerId == followeeId)
      return;
    if (const auto it = users.find(followerId);
        it != users.cend() && users.contains(followeeId))
      it->second.unfollow(followeeId);
  }

 private:
  int time = 0;
  unordered_map<int, User> users;  // {userId: User}
};

class Tweet {
  public int id;
  public int time;
  public Tweet next = null;
  public Tweet(int id, int time) {
    this.id = id;
    this.time = time;
  }
}

class User {
  private int id;
  public Set<Integer> followeeIds = new HashSet<>();
  public Tweet tweetHead = null;

  public User(int id) {
    this.id = id;
    follow(id);
  }

  public void follow(int followeeId) {
    followeeIds.add(followeeId);
  }

  public void unfollow(int followeeId) {
    followeeIds.remove(followeeId);
  }

  public void post(int tweetId, int time) {
    final Tweet oldTweetHead = tweetHead;
    tweetHead = new Tweet(tweetId, time);
    tweetHead.next = oldTweetHead;
  }
}

class Twitter {
  /** Compose a new tweet. */
  public void postTweet(int userId, int tweetId) {
    users.putIfAbsent(userId, new User(userId));
    users.get(userId).post(tweetId, time++);
  }

  /**
   * Retrieve 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
   * herself. Tweets must be ordered from most recent to least recent.
   */
  public List<Integer> getNewsFeed(int userId) {
    if (!users.containsKey(userId))
      return new ArrayList<>();

    List<Integer> newsFeed = new ArrayList<>();
    Queue<Tweet> maxHeap = new PriorityQueue<>((a, b) -> Integer.compare(b.time, a.time));

    for (final int followeeId : users.get(userId).followeeIds) {
      Tweet tweetHead = users.get(followeeId).tweetHead;
      if (tweetHead != null)
        maxHeap.offer(tweetHead);
    }

    int count = 0;
    while (!maxHeap.isEmpty() && count++ < 10) {
      Tweet tweet = maxHeap.poll();
      newsFeed.add(tweet.id);
      if (tweet.next != null)
        maxHeap.offer(tweet.next);
    }

    return newsFeed;
  }

  /**
   * Follower follows a followee.
   * If the operation is invalid, it should be a no-op.
   */
  public void follow(int followerId, int followeeId) {
    if (followerId == followeeId)
      return;
    users.putIfAbsent(followerId, new User(followerId));
    users.putIfAbsent(followeeId, new User(followeeId));
    users.get(followerId).follow(followeeId);
  }

  /**
   * Follower unfollows a followee.
   * If the operation is invalid, it should be a no-op.
   */
  public void unfollow(int followerId, int followeeId) {
    if (followerId == followeeId)
      return;
    if (users.containsKey(followerId) && users.containsKey(followeeId))
      users.get(followerId).unfollow(followeeId);
  }

  private int time = 0;
  private Map<Integer, User> users = new HashMap<>(); // {userId: User}
}

class Twitter:
  def __init__(self):
    self.timer = itertools.count(step=-1)
    self.tweets = collections.defaultdict(deque)
    self.followees = collections.defaultdict(set)

  def postTweet(self, userId: int, tweetId: int) -> None:
    self.tweets[userId].appendleft((next(self.timer), tweetId))
    if len(self.tweets[userId]) > 10:
      self.tweets[userId].pop()

  def getNewsFeed(self, userId: int) -> list[int]:
    tweets = list(
        heapq.merge(
            *
            (self.tweets[followee]
             for followee in self.followees[userId] | {userId})))
    return [tweetId for _, tweetId in tweets[:10]]

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

  def unfollow(self, followerId: int, followeeId: int) -> None:
    self.followees[followerId].discard(followeeId)