skiplist.go

package skiplist

import (
	// "encoding/json"
	// "fmt"
	"math/rand"
	"time"
)

const epsilon float64 = 0.0001

type SkipListNode struct {
	Member string
	Score  float64
	Next   []*SkipListNode
	Count  int
}

const LIST_DEPTH int = 5

type SkipList struct {
	Levels int           `json:"Levels"`
	Head   *SkipListNode `json:"Node,omitempty"`
}

func NewSkipList() *SkipList {
	sl := &SkipList{
		Levels: LIST_DEPTH,
		Head: &SkipListNode{
			Next: make([]*SkipListNode, LIST_DEPTH),
		},
	}
	rand.Seed(time.Now().UnixNano())
	return sl
}

type loopExitState int

const (
	loopContinue loopExitState = 0
	loopBreak    loopExitState = -2
	loopEnd      loopExitState = -3
)

func (sl *SkipList) checkState(curr, target *SkipListNode) loopExitState {
	if curr.Score == target.Score && curr.Member == target.Member {
		return loopEnd
	}

	if curr.Score > target.Score {
		return loopBreak
	}

	if curr.Score == target.Score && curr.Member > target.Member {
		return loopBreak
	}

	return loopContinue
}

func (sl *SkipList) Add(score float64, member string) {
	node := &SkipListNode{
		Member: member,
		Score:  score,
		Next:   make([]*SkipListNode, sl.Levels),
		Count:  1,
	}

	maxLevel := 1
	for maxLevel < sl.Levels && rand.Float64() < 0.5 {
		maxLevel++
	}

	cur := sl.Head

	// for each level.
	for i := maxLevel - 1; i >= 0; i-- {
		// loop and find the node where the new node is to be inserted
		// i.e traverse the elements as long as score is less
		count := cur.Count
		for cur.Next[i] != nil && sl.checkState(cur.Next[i], node) == loopContinue {
			cur = cur.Next[i]
			count += cur.Count
		}

		if cur.Next[i] != nil && sl.checkState(cur.Next[i], node) == loopEnd {
			return
		}

		node.Count = count
		node.Next[i] = cur.Next[i]
		cur.Next[i] = node
	}
}

type MemberWithScore struct {
	Score  float64
	Member interface{}
}

func (sl *SkipList) GetRange(minScore, maxScore float64) []*MemberWithScore {
	cur := sl.Head
	results := make([]*MemberWithScore, 0)

	for i := sl.Levels - 1; i >= 0; i++ {
		for cur.Next[i] != nil && cur.Next[i].Score < minScore {
			cur = cur.Next[i]
		}

		// traverse from the current level upto max maxScore

		cur = cur.Next[0]
		for cur != nil && cur.Score <= maxScore+epsilon {
			results = append(results, &MemberWithScore{Score: cur.Score, Member: cur.Member})
			cur = cur.Next[0]
		}

		if cur == nil {
			break
		}
	}

	return results
}

func (sl *SkipList) Count(minScore, maxScore float64) int {
	cur := sl.Head
	maxScore = maxScore + epsilon

	// b, _ := json.Marshal(sl)
	// fmt.Println(string(b))
	minNode := cur
	maxNode := cur
	for i := sl.Levels - 1; i >= 0; i-- {
		for minNode.Next[i] != nil && minNode.Next[i].Score < minScore {
			minNode = minNode.Next[i]
		}

		maxNode := minNode.Next[i]
		for maxNode != nil && maxNode.Next[i].Score < maxScore {
			maxNode = maxNode.Next[i]
		}
	}

	if maxNode == nil || minNode == nil {
		return -1
	}

	return maxNode.Count - minNode.Count
}

func (sl *SkipList) RemRangeByScore(minScore, maxScore float64) {
	cur := sl.Head
	for i := sl.Levels - 1; i >= 0; i-- {
		for cur.Next[i] != nil && cur.Next[i].Score < minScore {
			cur = cur.Next[i]
		}

		for cur.Next[i] != nil && cur.Next[i].Score <= maxScore+epsilon {
			cur.Next[i] = cur.Next[i].Next[i]
		}
	}
}