feat: progress

pkg/questions/design_twitter/main.go

@@ -0,0 +1,106 @@
+package design_twitter
+
+import (
+	"container/heap"
+	"maps"
+	"slices"
+)
+
+type TweetHeap [][]Tweet
+
+func (h TweetHeap) Len() int { return len(h) }
+func (h TweetHeap) Less(i, j int) bool {
+	timeI, timeJ := -1, -1
+
+	if len(h[i]) > 0 {
+		timeI = h[i][len(h[i])-1].Time
+	}
+	if len(h[j]) > 0 {
+		timeJ = h[j][len(h[i])-1].Time
+	}
+
+	return timeI > timeJ
+}
+func (h TweetHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
+
+func (h *TweetHeap) Push(x any) {
+	*h = append(*h, x.([]Tweet))
+}
+
+func (h *TweetHeap) Pop() any {
+	old := *h
+	x := old[len(old)-1]
+	*h = old[:len(old)-1]
+	return x
+}
+
+// ----------------------------------
+
+type Tweet struct {
+	ID   int
+	Time int
+}
+
+type Twitter struct {
+	tweets  map[int][]Tweet
+	follows map[int]map[int]bool
+	counter int
+}
+
+func Constructor() Twitter {
+	return Twitter{
+		tweets:  map[int][]Tweet{},
+		follows: map[int]map[int]bool{},
+		counter: 0,
+	}
+}
+
+func (t *Twitter) PostTweet(userId int, tweetId int) {
+	t.counter++
+
+	t.tweets[userId] = append(t.tweets[userId], Tweet{
+		ID:   tweetId,
+		Time: t.counter,
+	})
+}
+
+func (t *Twitter) GetNewsFeed(userId int) []int {
+	users := slices.AppendSeq([]int{userId}, maps.Keys(t.follows[userId]))
+
+	hp := TweetHeap{}
+	for _, user := range users {
+		hp = append(hp, t.tweets[user])
+	}
+	heap.Init(&hp)
+
+	feed := []int{}
+	for range 10 {
+		if len(hp[0]) == 0 {
+			break
+		}
+
+		top, rest := hp[0][len(hp[0])-1], hp[0][:len(hp[0])-1]
+		feed = append(feed, top.ID)
+		hp[0] = rest
+
+		heap.Fix(&hp, 0)
+	}
+
+	return feed
+}
+
+func (t *Twitter) Follow(followerId int, followeeId int) {
+	if _, ok := t.follows[followerId]; !ok {
+		t.follows[followerId] = map[int]bool{}
+	}
+
+	t.follows[followerId][followeeId] = true
+}
+
+func (t *Twitter) Unfollow(followerId int, followeeId int) {
+	if _, ok := t.follows[followerId]; !ok {
+		return
+	}
+
+	delete(t.follows[followerId], followeeId)
+}