设计朋友圈时间线功能

设计 Twitter:合并 k 个有序链表和面向对象设计

之所以要把 Tweet 和 User 类放到 Twitter 类里面,是因为 Tweet 类必须要用到一个全局时间戳 timestamp,而 User 类又需要用到 Tweet 类记录用户发送的推文,所以它们都作为内部类

#include <gtest/gtest.h>
#include <iostream>
#include <list>
#include <map>
#include <queue>
#include <set>

using namespace std;

class Twitter {
 private:
  class Tweet {
   public:
    int id;
    int time;     // 发表时间
    Tweet* next;  // 存在user链表,需要next指针

   public:
    Tweet(int id, int time) : id(id), time(time), next(NULL) {}
  };

  class User {
   public:
    int id;
    set<int> followed;  // 关注人列表
    Tweet* head;        // 该用户发表的tweet的头结点

   public:
    User(int id) : id(id), head(NULL) { follow(id); }

    void follow(int userID) { followed.insert(userID); }

    void unfollow(int userID) {
      if (userID != id) {
        followed.erase(userID);
      }
    }

    void post(int tweetID) {
      Tweet* t = new Tweet(tweetID, timestamp);
      timestamp++;
      t->next = head;
      head = t;
    }
  };

 public:
  // user 发表一条 tweet 动态
  void postTweet(int userID, int tweetID) {
    if (userMap.count(userID) == 0) {
      userMap[userID] = new User(userID);
    }
    userMap[userID]->post(tweetID);
  }

  // 返回该user关注的人(包括自己)最近的动态ID,最多10条,这些动态按从新到旧的时间线排序
  // 合并多个有序链表
  list<int> getNewsFeed(int userID) {
    list<int> res;
    if (userMap.count(userID) == 0) return res;

    set<int> followees = userMap[userID]->followed;
    auto comp = [](Tweet* t1, Tweet* t2) { return t1->time > t2->time; };
    priority_queue<Tweet*, std::vector<Tweet*>, decltype(comp)> pq(comp);

    for (auto& user : followees) {
      auto t = userMap[user]->head;
      if (!t) continue;
      pq.push(t);
    }

    while (!pq.empty()) {
      if (res.size() == 10) break;
      auto t = pq.top();
      res.push_back(t->id);
      pq.pop();
      if (t->next) pq.push(t->next);
    }

    return res;
  }

  // follower 关注 followee, 如果ID不存在则新建
  void follow(int followerID, int followeeID) {
    if (userMap.count(followerID) == 0) {
      userMap[followerID] = new User(followerID);
    }
    if (userMap.count(followeeID) == 0) {
      userMap[followeeID] = new User(followeeID);
    }
    userMap[followerID]->follow(followeeID);
  }

  // follower 取关 followee, 如果ID不存在则什么都不做
  void unfollow(int followerID, int followeeID) {
    if (userMap.count(followerID) != 0) {
      userMap[followerID]->unfollow(followeeID);
    }
  }

 private:
  map<int, User*> userMap;
  static int timestamp;
};

int Twitter::timestamp = 0;

void printList(list<int> l) {
  for (auto v : l) {
    cout << v << " ";
  }
  cout << endl;
}

TEST(Algo, twitter) {
  Twitter tw;
  tw.postTweet(1, 5);
  printList(tw.getNewsFeed(1));
  tw.follow(1, 2);
  tw.postTweet(2, 6);
  printList(tw.getNewsFeed(1));
  tw.unfollow(1, 2);
  printList(tw.getNewsFeed(1));
}