利用 Java 8 的特性,快速简洁的实现各个操作
class Twitter {
Map<Integer, HashSet<Integer>> followMap = new HashMap<>();
Map<Integer, List<Pair>> tweetMap = new HashMap<>();
int ts;
public Twitter() {
ts = 0;
}
public void postTweet(int userId, int tweetId) {
this.tweetMap.computeIfAbsent(userId, k -> new ArrayList<>()).add(new Pair(ts, tweetId));
ts++;
}
public List<Integer> getNewsFeed(int userId) {
// 利用多路归并+大顶堆来获取前 10 条
// 利用时间戳排序
PriorityQueue<Integer[]> heap = new PriorityQueue<>((o1, o2) -> o2[0] - o1[0]);
// 插入自己。set 会去重
this.followMap.computeIfAbsent(userId, k -> new HashSet<>()).add(userId);
for (Integer user : this.followMap.get(userId)) {
List<Pair> p = this.tweetMap.get(user);
if (p != null && !p.isEmpty()) {
int i = p.size() - 1;
heap.add(new Integer[]{p.get(i).ts, p.get(i).tid, i, user});
}
}
List<Integer> res = new ArrayList<>();
for (int i = 0; i < 10 && !heap.isEmpty(); i++) {
Integer[] top = heap.poll();
res.add(top[1]);
int j = top[2];
// 多路归并
if (j > 0) {
j--;
Integer user = top[3];
List<Pair> p = this.tweetMap.get(user);
heap.add(new Integer[]{p.get(j).ts, p.get(j).tid, j, user});
}
}
return res;
}
public void follow(int followerId, int followeeId) {
this.followMap.computeIfAbsent(followerId, k -> new HashSet<>()).add(followeeId);
}
public void unfollow(int followerId, int followeeId) {
this.followMap.computeIfPresent(followerId, (k, v) -> {
v.remove(followeeId);
if (v.isEmpty()) return null;
return v;
});
}
}
class Pair {
int ts, tid;
Pair(int ts, int tid) {
this.ts = ts;
this.tid = tid;
}
}
/**
* Your Twitter object will be instantiated and called as such:
* Twitter obj = new Twitter();
* obj.postTweet(userId,tweetId);
* List<Integer> param_2 = obj.getNewsFeed(userId);
* obj.follow(followerId,followeeId);
* obj.unfollow(followerId,followeeId);
*/