Zenn
🐙

NeetCode 150 [Heap / Priority Queue]medium:Design Twitter

に公開
heap
neetcode
priorityqueue
tech

NeetCodeのSolutionを書いていく

Design Twitter

問題

簡単なTwitterを作りましょう。
ユーザはpostができ、お互いにfollow, unfollowすることもできます。
そして自分自身の登録した10個の最新のtweetを見ることができます。

ユーザとtweetはユニークなID(整数)で区別できます。
次の関数を実装しましょう。

  • Twitter() twitterオブジェクトを初期化する
  • void postTweet(int usrId, int tweetId) userIdのユーザでtweetIdのtweetを発行します。tweetIdはユニークです。
  • List<Integer> getNewsFeed(int userId) ユーザの登録しているニュースから最大で最新10個のtweetを取得します。それぞれのtweetはユーザ自身か、ユーザがfollowしているユーザによるpostです。Tweet IDは降順である必要があります。
  • void follow(int followerId, int followeeId) followerIdのユーザがfolloweeIdのユーザをfollowします。
  • void unfollow(int followerId, int followeeId) followerIdのユーザがfolloweeIdのユーザをunfollowします。

例

Input:
["Twitter", "postTweet", [1, 10], "postTweet", [2, 20], "getNewsFeed", [1], "getNewsFeed", [2], "follow", [1, 2], "getNewsFeed", [1], "getNewsFeed", [2], "unfollow", [1, 2], "getNewsFeed", [1]]

Output:
[null, null, null, [10], [20], null, [20, 10], [20], null, [10]]

Explanation:
Twitter twitter = new Twitter();
twitter.postTweet(1, 10); // User 1 posts a new tweet with id = 10.
twitter.postTweet(2, 20); // User 2 posts a new tweet with id = 20.
twitter.getNewsFeed(1);   // User 1's news feed should only contain their own tweets -> [10].
twitter.getNewsFeed(2);   // User 2's news feed should only contain their own tweets -> [20].
twitter.follow(1, 2);     // User 1 follows user 2.
twitter.getNewsFeed(1);   // User 1's news feed should contain both tweets from user 1 and user 2 -> [20, 10].
twitter.getNewsFeed(2);   // User 2's news feed should still only contain their own tweets -> [20].
twitter.unfollow(1, 2);   // User 1 follows user 2.
twitter.getNewsFeed(1);   // User 1's news feed should only contain their own tweets -> [10].

制約

  • 1 <= userId, followerId, followeeId <= 100
  • 0 <= tweetId <= 1000
  • 計算量；getNewsFeed()関数はO(nlogn)を目指してください。それ以外はO(1)です。
  • メモリ：O(N * m) + (N * M) + n
    nはuserIdに紐づくfolloweeIdの数、mはどのユーザにかによらない最大のtweet数。NはuserIdの数、Mはどのユーザかに依らない最大のfollowee数。

メモ

コンストラクタ１つ、関数４つ。
時間がかかりそうな問題。
でもこういう系は考えることが少なくて簡単なはず。

getNewsFeed以外はO(1)である。
getNewsFeedで返却するのは最大10個であり、ユーザ数分だけ保持する必要がある？

userId、tweetId両方ユニークであることを考えるとdictを使えばよいか？
postTweetで以下のdictをqueueに詰める。
postInfo :{ userId : [tweetId] }
follow, unfollowで以下のdictを操作
followerInfo : { follower : [followee] }
getNewsFeedでfollwerのfolloweeのIdと自分自身のIdのpostInfoの情報を取得。
heapqで10個目までを返す？
この時投稿の番号を覚えておくために投稿と番号をくっつけておく必要がある

import heapq
from typing import DefaultDict


class Twitter:
    def __init__(self):
        self.postInfo = DefaultDict(list)
        self.followerInfo = DefaultDict(set)
        self.count = 0
        pass

    def postTweet(self, userId: int, tweetId: int) -> None:
        self.count -= 1
        self.postInfo[userId].append([self.count, tweetId])

    def getNewsFeed(self, userId: int) -> list[int]:
        followers = self.followerInfo.get(userId, set())
        if userId not in followers:
            followers.add(userId)

        posts = []
        for follower in followers:
            if self.postInfo.get(follower):
                for count, tweetId in self.postInfo[follower]:
                    heapq.heappush(posts, [count, tweetId])

        top_posts = []
        for _ in range(min(10, len(posts))):
            top_posts.append(heapq.heappop(posts)[1])
        return top_posts

    def follow(self, followerId: int, followeeId: int) -> None:
        if self.followerInfo.get(followerId):
            self.followerInfo[followerId].add(followeeId)
        else:
            self.followerInfo[followerId] = {followeeId}

    def unfollow(self, followerId: int, followeeId: int) -> None:
        if followeeId in self.followerInfo[followerId]:
            self.followerInfo[followerId].remove(followeeId)

Discussion