github.com/cloudcentricdev/golang-tutorials/03@v0.0.0-20231018054503-b24024842337/skiplist/skiplist.go (about)

     1  package skiplist
     2  
     3  import (
     4  	"bytes"
     5  	"errors"
     6  	"math"
     7  
     8  	"github.com/cloudcentricdev/golang-tutorials/03/fastrand"
     9  )
    10  
    11  const (
    12  	MaxHeight = 16
    13  	PValue    = 0.5 // p = 1/2
    14  )
    15  
    16  var probabilities [MaxHeight]uint32
    17  
    18  type node struct {
    19  	key   []byte
    20  	val   []byte
    21  	tower [MaxHeight]*node
    22  }
    23  
    24  type SkipList struct {
    25  	head   *node
    26  	height int
    27  }
    28  
    29  func NewSkipList() *SkipList {
    30  	sl := &SkipList{}
    31  	sl.head = &node{}
    32  	sl.height = 1
    33  	return sl
    34  }
    35  
    36  func init() {
    37  	probability := 1.0
    38  
    39  	for level := 0; level < MaxHeight; level++ {
    40  		probabilities[level] = uint32(probability * float64(math.MaxUint32))
    41  		probability *= PValue
    42  	}
    43  }
    44  
    45  func randomHeight() int {
    46  	seed := fastrand.Uint32()
    47  
    48  	height := 1
    49  	for height < MaxHeight && seed <= probabilities[height] {
    50  		height++
    51  	}
    52  
    53  	return height
    54  }
    55  
    56  func (sl *SkipList) search(key []byte) (*node, [MaxHeight]*node) {
    57  	var next *node
    58  	var journey [MaxHeight]*node
    59  
    60  	prev := sl.head
    61  	for level := sl.height - 1; level >= 0; level-- {
    62  		for next = prev.tower[level]; next != nil; next = prev.tower[level] {
    63  			if bytes.Compare(key, next.key) <= 0 {
    64  				break
    65  			}
    66  			prev = next
    67  		}
    68  		journey[level] = prev
    69  	}
    70  
    71  	if next != nil && bytes.Equal(key, next.key) {
    72  		return next, journey
    73  	}
    74  	return nil, journey
    75  }
    76  
    77  func (sl *SkipList) Find(key []byte) ([]byte, error) {
    78  	found, _ := sl.search(key)
    79  
    80  	if found == nil {
    81  		return nil, errors.New("key not found")
    82  	}
    83  
    84  	return found.val, nil
    85  }
    86  
    87  func (sl *SkipList) Insert(key []byte, val []byte) {
    88  	found, journey := sl.search(key)
    89  
    90  	if found != nil {
    91  		// update value of existing key
    92  		found.val = val
    93  		return
    94  	}
    95  	height := randomHeight()
    96  	nd := &node{key: key, val: val}
    97  
    98  	for level := 0; level < height; level++ {
    99  		prev := journey[level]
   100  
   101  		if prev == nil {
   102  			// prev is nil if we are extending the height of the tree,
   103  			// because that level did not exist while the journey was being recorded
   104  			prev = sl.head
   105  		}
   106  		nd.tower[level] = prev.tower[level]
   107  		prev.tower[level] = nd
   108  	}
   109  
   110  	if height > sl.height {
   111  		sl.height = height
   112  	}
   113  }
   114  
   115  func (sl *SkipList) Delete(key []byte) bool {
   116  	found, journey := sl.search(key)
   117  
   118  	if found == nil {
   119  		return false
   120  	}
   121  
   122  	for level := 0; level < sl.height; level++ {
   123  		if journey[level].tower[level] != found {
   124  			break
   125  		}
   126  		journey[level].tower[level] = found.tower[level]
   127  		found.tower[level] = nil
   128  	}
   129  	found = nil
   130  	sl.shrink()
   131  
   132  	return true
   133  }
   134  
   135  func (sl *SkipList) shrink() {
   136  	for level := sl.height - 1; level >= 0; level-- {
   137  		if sl.head.tower[level] == nil {
   138  			sl.height--
   139  		}
   140  	}
   141  }
   142  
   143  func (sl *SkipList) String() string {
   144  	v := &visualizer{sl}
   145  	return v.visualize()
   146  }
   147  
   148  type Iterator struct {
   149  	current *node
   150  }
   151  
   152  func (i *Iterator) hasNext() bool {
   153  	return i.current.tower[0] != nil
   154  }
   155  
   156  func (i *Iterator) next() ([]byte, []byte) {
   157  	i.current = i.current.tower[0]
   158  
   159  	if i.current == nil {
   160  		return nil, nil
   161  	}
   162  	return i.current.key, i.current.val
   163  }
   164  
   165  func (sl *SkipList) Iterator() *Iterator {
   166  	return &Iterator{sl.head.tower[0]}
   167  }