go-common/app/job/main/up-rating/service/trend_heap.go
2019-04-22 18:49:16 +08:00

102 lines
1.8 KiB
Go

package service
import (
"container/heap"
"go-common/app/job/main/up-rating/model"
)
// Heap for diff topK
type Heap interface {
Put(*model.Diff)
Result() []*model.Diff
}
// AscHeap for asc
type AscHeap struct {
heap []*model.Diff
ctype int
}
// Put put diff to heap
func (a *AscHeap) Put(diff *model.Diff) {
heap.Push(a, diff)
if a.Len() > 200 {
heap.Pop(a)
}
}
// Result get result
func (a *AscHeap) Result() []*model.Diff {
return a.heap
}
// Len len
func (a *AscHeap) Len() int { return len(a.heap) }
// Less less
func (a *AscHeap) Less(i, j int) bool {
return a.heap[i].GetScore(a.ctype) < a.heap[j].GetScore(a.ctype)
}
// Swap swap
func (a *AscHeap) Swap(i, j int) { a.heap[i], a.heap[j] = a.heap[j], a.heap[i] }
// Push push to heap
func (a *AscHeap) Push(x interface{}) {
a.heap = append(a.heap, x.(*model.Diff))
}
// Pop pop from heap
func (a *AscHeap) Pop() interface{} {
old := a.heap
n := len(old)
x := old[n-1]
a.heap = old[0 : n-1]
return x
}
// DescHeap for desc
type DescHeap struct {
heap []*model.Diff
ctype int
}
// Put to descHeap
func (d *DescHeap) Put(diff *model.Diff) {
heap.Push(d, diff)
if d.Len() > 200 {
heap.Pop(d)
}
}
// Result desc heap result
func (d *DescHeap) Result() []*model.Diff {
return d.heap
}
// Len len
func (d *DescHeap) Len() int { return len(d.heap) }
// Less less
func (d *DescHeap) Less(i, j int) bool {
return d.heap[i].GetScore(d.ctype) > d.heap[j].GetScore(d.ctype)
}
// Swap swap
func (d *DescHeap) Swap(i, j int) { d.heap[i], d.heap[j] = d.heap[j], d.heap[i] }
// Push push to desc heap
func (d *DescHeap) Push(x interface{}) {
d.heap = append(d.heap, x.(*model.Diff))
}
// Pop pop from desc heap
func (d *DescHeap) Pop() interface{} {
old := d.heap
n := len(old)
x := old[n-1]
d.heap = old[0 : n-1]
return x
}