107 lines
1.8 KiB
Go
107 lines
1.8 KiB
Go
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)
|
|
}
|