leetCode：Design Twitter

leetCode の「Design Twitter」を解いてみた。この問題は、簡易版の Twiter を作成するというもの。
https://leetcode.com/problems/design-twitter

実装する上で気にしたのは以下の点。

• ユーザやフォロワーの検索を速くするため、これらの入れ物には ArrayList ではなく HashSet や HashMap を使用している。
• ツイートを時刻でソートして取り出す処理を高速化するため、PriorityQueue を使用している。
• 各ツイートに持たせる時刻は、ミリ秒単位ではなく、ナノ秒単位にしている。これはミリ秒単位だと、連続ツイートの際に時刻が全く同じになってしまうことがあり、時刻でのソートがうまく効かないため。
• 本来であれば、ユーザ登録が最初に来ると思うのだが、この問題ではユーザ登録機能はないので、新しいユーザが出てくるたびにユーザを作成し、ユーザ一覧のMapに登録している。

提出したコードは以下。

import java.util.*;

public class Twitter {
    private Map<Integer, User> userIdToUser;

    /** Initialize your data structure here. */
    public Twitter() {
        userIdToUser = new HashMap<>();
    }

    /** Compose a new tweet. */
    public void postTweet(int userId, int tweetId) {
        Tweet tweet = new Tweet(tweetId, System.nanoTime());

        User user = userIdToUser.get(userId);
        if (user == null) {
            User newUser = registerUser(userId);
            newUser.getTweets().add(tweet);
        }
        else {
            userIdToUser.get(userId).getTweets().add(tweet);
        }
    }

    /** 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) {
        List<Integer> newsFeedIds = new ArrayList<>();
        User user = userIdToUser.get(userId);
        if (user == null) {
            return newsFeedIds;
        }

        Queue<Tweet> tweetQueue = new PriorityQueue(Comparator.comparing(Tweet::getTime, Comparator.reverseOrder()));
        for (Tweet tweet : user.getTweets()) {
            tweetQueue.add(tweet);
        }
        for (User follower : user.getFollowers()) {
            if (follower.equals(user)) {
                continue;
            }
            for (Tweet tweet : follower.getTweets()) {
                tweetQueue.add(tweet);
            }
        }

        int count = 0;
        while(!tweetQueue.isEmpty() && count < 10) {
            newsFeedIds.add(tweetQueue.poll().getTweetId());
            count++;
        }
        return newsFeedIds;
    }

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

        User  followee = userIdToUser.get(followeeId);
        if (followee == null) {
            followee = registerUser(followeeId);
        }

        follower.getFollowers().add(followee);
    }

    /** Follower unfollows a followee. If the operation is invalid, it should be a no-op. */
    public void unfollow(int followerId, int followeeId) {
        User follower = userIdToUser.get(followerId);
        if (follower == null) {
            follower = registerUser(followerId);
        }

        User  followee = userIdToUser.get(followeeId);
        if (followee == null) {
            followee = registerUser(followeeId);
        }

        follower.getFollowers().remove(followee);
    }

    private User registerUser(int userId) {
        User newUser = new User();
        userIdToUser.put(userId, newUser);
        return newUser;
    }

    private static class User {
        List<Tweet> tweets = new ArrayList<>();
        Set<User> followers = new HashSet<>();

        private List<Tweet> getTweets() {
            return tweets;
        }

        private Set<User> getFollowers() {
            return followers;
        }
    }

    private static class Tweet {
        private int tweetId;
        private long time;

        private Tweet(int tweetId, long time) {
            this.tweetId = tweetId;
            this.time = time;
        }

        private int getTweetId() {
            return tweetId;
        }

        private long getTime() {
            return time;
        }
    }
}

上記を提出したら、

Runtime: 29 ms, faster than 49.75% of Java online submissions for Design Twitter.
Memory Usage: 46.9 MB, less than 100.00% of Java online submissions for Design Twitter.

となった。提出者の中では真ん中くらいのスピード。メモリの使用量はかなり少なく済んでいる。

上述のコードのうち getNewsFeed メソッドについては、まだまだ高速化の余地はあると思うので、時間のある時にさらなる高速化を試みたい。