Design Twitter

Medium

The problem

Build a tiny Twitter. Users can PostTweet, Follow and Unfollow other users, and GetNewsFeed(userId) returns the ids of the 10 most recent tweets from the user and the people they follow, newest first. Fewer than 10 are returned if fewer exist.

Example 1
Input: PostTweet(1, 5) GetNewsFeed(1) Follow(1, 2) PostTweet(2, 6) GetNewsFeed(1) Unfollow(1, 2) GetNewsFeed(1)
Output: [5], [6, 5], [5]
User 1 sees their own tweet 5. After following user 2, the newer tweet 6 comes first. After unfollowing, only 5 remains. (Other calls return nothing.)

Limits

1 ≤ userId, followerId, followeeId ≤ 500
0 ≤ tweetId ≤ 10,000, all tweet ids are different
Up to 30,000 calls in total
A user does not need to follow themselves, and their own tweets always show

Write it in Go. Try for about 20 minutes on paper first, then open one hint at a time.

Try it here

Write Go. Common packages like fmt and sort are imported for you. Keep the function name and inputs the same.

Hints, one at a time

Nudge

getNewsFeed is "merge k sorted lists" in disguise — each followee has a list of tweets by time.

The idea

Store tweets per user with a global timestamp and a follow set. For the feed, push each followee's latest tweet into a max-heap by time and pop 10, pushing the previous tweet of that user.

Target: O(k + 10 log k) per feed

Go function shape

type Twitter struct{}
func Constructor() Twitter
func (t *Twitter) PostTweet(userId int, tweetId int)
func (t *Twitter) GetNewsFeed(userId int) []int
func (t *Twitter) Follow(followerId int, followeeId int)
func (t *Twitter) Unfollow(followerId int, followeeId int)

Reference solution

Tested with go test. Try it yourself first, then compare.

// Twitter keeps each user's tweets (newest last) and follow set. The feed is "merge k sorted lists":
// a max-heap on tweet time holds each followed user's newest tweet, and we pop 10, refilling from the same user.
type Twitter struct {
	clock   int
	tweets  map[int][]tweet
	follows map[int]map[int]bool
}
type tweet struct{ time, id int }

func NewTwitter() *Twitter {
	return &Twitter{tweets: map[int][]tweet{}, follows: map[int]map[int]bool{}}
}

func (t *Twitter) PostTweet(user, id int) {
	t.clock++
	t.tweets[user] = append(t.tweets[user], tweet{t.clock, id})
}

func (t *Twitter) Follow(follower, followee int) {
	if t.follows[follower] == nil {
		t.follows[follower] = map[int]bool{}
	}
	t.follows[follower][followee] = true
}

func (t *Twitter) Unfollow(follower, followee int) { delete(t.follows[follower], followee) }

type feedItem struct {
	tw    tweet
	user  int
	index int // position of tw in that user's list, so the previous tweet is index-1
}
type feedHeap []feedItem

func (h feedHeap) Len() int           { return len(h) }
func (h feedHeap) Less(i, j int) bool { return h[i].tw.time > h[j].tw.time }
func (h feedHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *feedHeap) Push(x any)        { *h = append(*h, x.(feedItem)) }
func (h *feedHeap) Pop() any {
	old := *h
	x := old[len(old)-1]
	*h = old[:len(old)-1]
	return x
}

func (t *Twitter) GetNewsFeed(user int) []int {
	h := &feedHeap{}
	add := func(u int) {
		if list := t.tweets[u]; len(list) > 0 {
			heap.Push(h, feedItem{list[len(list)-1], u, len(list) - 1})
		}
	}
	add(user) // you see your own tweets too
	for f := range t.follows[user] {
		if f != user {
			add(f)
		}
	}
	feed := []int{}
	for h.Len() > 0 && len(feed) < 10 {
		top := heap.Pop(h).(feedItem)
		feed = append(feed, top.tw.id)
		if top.index > 0 {
			heap.Push(h, feedItem{t.tweets[top.user][top.index-1], top.user, top.index - 1})
		}
	}
	return feed
}

Read the lesson: Heap / Top-K