package zset
import ("math/rand")const(
skipListMaxLevel= 8 //(1/p)^maxLevel >= maxNode
skipListP = 0.25 //SkipList P = 1/4
)//Element e
type Element struct{
time int64
key uint64
}//Key return key
func (e *Element) Key() uint64 {returne.key
}//Time 时间
func (e *Element) Time() int64 {returne.time
}
type zSkipListLevelstruct{
forward*ZSkipListNode
span uint32
}//ZSkipListNode is an element of a skip list
type ZSkipListNode struct{
ele*Element
score uint32
backward*ZSkipListNode
level []zSkipListLevel
orderint}
func zslCreateNode(levelint, score uint32, ele *Element) *ZSkipListNode {
zn := &ZSkipListNode{
ele: ele,
score: score,
level: make([]zSkipListLevel, level),
}returnzn
}//Score return score
func (node *ZSkipListNode) Score() uint32 {returnnode.score
}//Element return Element
func (node *ZSkipListNode) Element() *Element {returnnode.ele
}//zSkipList represents a skip list
type zSkipList struct{
header, tail*ZSkipListNode
length uint32
levelint //current level count
}//zslCreate creates a skip list
func zslCreate() *zSkipList {
zsl := &zSkipList{
level:1,
}
zsl.header= zslCreateNode(skipListMaxLevel, 0, nil)returnzsl
}//insert element
func (list *zSkipList) insert(node *ZSkipListNode) *ZSkipListNode {var update [skipListMaxLevel]*ZSkipListNodevarrank [skipListMaxLevel]uint32
x :=list.headerfor i := list.level - 1; i >= 0; i--{if i == list.level-1{
rank[i]= 0}else{
rank[i]= rank[i+1]
}for x.level[i].forward != nil &&(x.level[i].forward.score< node.score ||x.level[i].forward.score== node.score &&node.ele.Time()
rank[i]+=x.level[i].span
x=x.level[i].forward
}
update[i]=x
}
level :=len(node.level)if level >list.level {for i := list.level; i < level; i++{
rank[i]= 0update[i]=list.header
update[i].level[i].span=list.length
}
list.level=level
}
x=nodefor i := 0; i < level; i++{
x.level[i].forward=update[i].level[i].forward
update[i].level[i].forward=x
x.level[i].span= update[i].level[i].span - (rank[0] -rank[i])
update[i].level[i].span= (rank[0] - rank[i]) + 1}for i := level; i < list.level; i++{
update[i].level[i].span++}
next := x.level[0].forwardif next != nil && x.score == next.score && x.ele.Time() ==next.ele.Time() {
x.order= next.order + 1}if update[0] ==list.header {
x.backward=nil
}else{
x.backward= update[0]
}if x.level[0].forward ==nil {
list.tail=x
}else{
x.level[0].forward.backward =x
}
list.length++
returnx
}//delete element
func (list *zSkipList) delete(node *ZSkipListNode) *ZSkipListNode {var update [skipListMaxLevel]*ZSkipListNode
x :=list.headerfor i := list.level - 1; i >= 0; i--{for next := x.level[i].forward; next != nil &&(next.score< node.score ||next.score== node.score &&(node.ele.Time()< next.ele.Time() ||node.ele.Time()== next.ele.Time() && node.order < next.order)); next =x.level[i].forward {
x=next
}
update[i]=x
}
x= x.level[0].forwardif x != nil && x.score == node.score && x.ele.key ==node.ele.key {for i := 0; i < list.level; i++{if update[i].level[i].forward ==x {
update[i].level[i].span+= x.level[i].span - 1update[i].level[i].forward=x.level[i].forward
}else{
update[i].level[i].span--}
}if x.level[0].forward ==nil {
list.tail=x.backward
}else{
x.level[0].forward.backward =x.backward
}for list.level > 1 && list.header.level[list.level-1].forward ==nil {
list.level--}
list.length--
returnx
}returnnil
}//Find the rank for an element.//Returns 0 when the element cannot be found, rank otherwise.//Note that the rank is 1-based
func (list *zSkipList) zslGetRank(node *ZSkipListNode) uint32 {varrank uint32
x :=list.headerfor i := list.level - 1; i >= 0; i--{for next := x.level[i].forward; next != nil &&(next.score< node.score ||next.score== node.score &&(node.ele.time< next.ele.time ||node.ele.time== next.ele.time && node.order <= next.order)); next =x.level[i].forward {
rank+=x.level[i].span
x=next
}if x.ele != nil && x.ele.key ==node.ele.key {returnrank
}
}return 0}
func (list*zSkipList) randomLevel() int{
lvl := 1
for lvl < skipListMaxLevel && rand.Float64()
lvl++}returnlvl
}//Finds an element by its rank. The rank argument needs to be 1-based.
func (list *zSkipList) getElementByRank(rank uint32) *ZSkipListNode {if rank ==list.length {returnlist.tail
}if rank == 1{return list.header.level[0].forward
}vartraversed uint32
x :=list.headerfor i := list.level - 1; i >= 0; i--{for x.level[i].forward != nil && traversed+x.level[i].span <=rank {
traversed+=x.level[i].span
x=x.level[i].forward
}if traversed ==rank {returnx
}
}returnnil
}//ZSet set
type ZSet struct{
dict map[uint64]*ZSkipListNode
zsl*zSkipList
}//NewZSet create ZSet
func NewZSet() *ZSet {
zs := &ZSet{
dict: make(map[uint64]*ZSkipListNode),
zsl: zslCreate(),
}returnzs
}//Add a new element or update the score of an existing element
func (zs *ZSet) Add(score uint32, key uint64, t int64) *ZSkipListNode {if node := zs.dict[key]; node !=nil {
oldScore :=node.scoreif score ==oldScore {returnnil
}if next := node.level[0].forward; score > oldScore && (next == nil || score
node.score=score
node.ele.time=t
}else if score < oldScore && (node.backward == nil || score >node.backward.score) {
node.score=score
node.ele.time=t
}else{
zs.zsl.delete(node)
node.score=score
node.ele.time=t
zs.zsl.insert(node)
}returnnode
}else{
ele := &Element{
key: key,
time: t,
}
lvl :=zs.zsl.randomLevel()
node :=zslCreateNode(lvl, score, ele)
zs.zsl.insert(node)
zs.dict[key]=nodereturnnode
}
}//Delete the element 'ele' from the sorted set,//return 1 if the element existed and was deleted, 0 otherwise
func (zs *ZSet) Delete(id uint64) int{
node :=zs.dict[id]if node ==nil {return 0}
zs.zsl.delete(node)
delete(zs.dict, id)return 1}//Rank return 1-based rank or 0 if not exist
func (zs *ZSet) Rank(id uint64, reverse bool) (uint32, uint32) {
node :=zs.dict[id]if node !=nil {
rank :=zs.zsl.zslGetRank(node)if rank > 0{ifreverse {return zs.zsl.length - rank + 1, node.score
}returnrank, node.score
}
}return 0, 0}//Score return score
func (zs *ZSet) Score(id uint64) uint32 {
node :=zs.dict[id]if node !=nil {returnnode.score
}return 0}//Range return 1-based elements in [start, end]
func (zs *ZSet) Range(start uint32, end uint32, reverse bool, retKey *[]uint64, retScore *[]uint32) {if start == 0{
start= 1}if end == 0{
end=zs.zsl.length
}if start > end || start >zs.zsl.length {return}if end >zs.zsl.length {
end=zs.zsl.length
}
rangeLen := end - start + 1
ifreverse {
node := zs.zsl.getElementByRank(zs.zsl.length - start + 1)for i := uint32(0); i < rangeLen; i++{*retKey = append(*retKey, node.ele.key)*retScore = append(*retScore, node.score)
node=node.backward
}
}else{
node :=zs.zsl.getElementByRank(start)for i := uint32(0); i < rangeLen; i++{*retKey = append(*retKey, node.ele.key)*retScore = append(*retScore, node.score)
node= node.level[0].forward
}
}
}//Length return the element count
func (zs *ZSet) Length() uint32 {returnzs.zsl.length
}//MinScore return min score
func (zs *ZSet) MinScore() uint32 {
first := zs.zsl.header.level[0].forwardif first !=nil {returnfirst.score
}return 0}//Tail return the last element
func (zs *ZSet) Tail() *Element {if zs.zsl.tail !=nil {returnzs.zsl.tail.ele
}returnnil
}//DeleteFirst the first element
func (zs *ZSet) DeleteFirst() *Element {
node := zs.zsl.header.level[0].forward
zs.zsl.delete(node)
delete(zs.dict, node.ele.key)returnnode.ele
}