数据结构模板

1. 跳跃表（skiplist）

一种简单的 skiplist 实现，关于跳跃表的原理可以参考其它文章。

package skiplist

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

const (
	maxLevel = 16
)

var (
	random *rand.Rand
)

// SLKey skip_list key 定义
type SLKey interface {
	Less(aKey SLKey) bool
	Equal(aKey SLKey) bool
}

// IntKey 整数key
type IntKey int

// Less 整数Key比较函数
func (ik IntKey) Less(aKey SLKey) bool {
	if _, ok := aKey.(IntKey); !ok {
		return false
	}
	ik2, _ := aKey.(IntKey)
	return ik < ik2
}

// Equal key相等比较
func (ik IntKey) Equal(aKey SLKey) bool {
	if _, ok := aKey.(IntKey); !ok {
		return false
	}
	ik2, _ := aKey.(IntKey)
	return ik == ik2
}

func max(a int, b int) int {
	if a >= b {
		return a
	}
	return b
}

// SkipList 跳跃表struct定义
type SkipList struct {
	level int   // 当前跳跃表的最大层数
	head  *Node // 头节点指针
}

// Node 节点struct定义
type Node struct {
	key  SLKey
	val  interface{}
	next []*Node
}

// 初始化
func init() {
	random = rand.New(rand.NewSource(time.Now().UnixNano()))
}

// NewNode 新建节点
func NewNode(key SLKey, val interface{}, n int) *Node {
	return &Node{
		key:  key,
		val:  val,
		next: make([]*Node, n), // 创建maxLevel
	}
}

// NewSkipList 新建跳跃表节点
func NewSkipList() *SkipList {
	return &SkipList{
		level: 0,
		head:  NewNode(IntKey(0), 0, maxLevel),
	}
}

// level取值范围 0 ~ maxLevel - 1
func getInsertLevel() int {
	level := 0
	for random.Intn(2) == 1 && level < maxLevel-1 {
		level++
	}
	return level
}

// SearchKey 查找key对应的val
func (sl *SkipList) SearchKey(aKey SLKey) interface{} {

	if sl.head.next[0] == nil {
		return nil
	}

	cl := sl.level
	p := sl.head
	cn := sl.head
	for cl >= 0 {
		for cn = cn.next[cl]; cn != nil && cn.key.Less(aKey); p, cn = cn, cn.next[cl] {
		}
		if cn != nil && cn.key.Equal(aKey) {
			return cn.val
		}
		cn = p
		cl--
	}
	return nil
}

// InsertKey 插入函数, 如果key已经存在，则更新val
func (sl *SkipList) InsertKey(aKey SLKey, val interface{}) {
	il := getInsertLevel()
	cl := max(il, sl.level)
	upNds := make([]*Node, cl+1)
	p := sl.head
	cn := sl.head
	var snd *Node
	for cl >= 0 {
		for cn = cn.next[cl]; cn != nil && cn.key.Less(aKey); p, cn = cn, cn.next[cl] {
		}
		if cn != nil && cn.key.Equal(aKey) {
			snd = cn
			cn.val = val
		} else {
			upNds[cl] = p
		}
		cn = p
		cl--
	}

	// 创建新对象
	if snd == nil {
		snd = NewNode(aKey, val, maxLevel)
	}

	// 更新跳跃链
	for i := il; i >= 0; i-- {
		if upNds[i] != nil {
			snd.next[i] = upNds[i].next[i]
			upNds[i].next[i] = snd
		}
	}

	sl.level = max(sl.level, il)
	return
}

// DeleteKey 删除key
func (sl *SkipList) DeleteKey(aKey SLKey) {
	cl := sl.level
	p := sl.head
	cn := sl.head
	upNds := make([]*Node, sl.level+1)
	var snd *Node

	for cl >= 0 {
		for cn = cn.next[cl]; cn != nil && cn.key.Less(aKey); p, cn = cn, cn.next[cl] {
		}
		if cn != nil && cn.key.Equal(aKey) {
			snd = cn
			upNds[cl] = p
		}
		cn = p
		cl--
	}

	// 不存在，直接返回
	if snd == nil {
		return
	}

	for i := sl.level; i >= 0; i-- {
		if upNds[i] != nil {
			upNds[i].next[i] = snd.next[i] // 直接干掉snd即可
		}
		if sl.head.next[i] == nil && i > 0 {
			sl.level--
		}
	}

	return
}

// PrintElem 打印所有层级元素
func (sl *SkipList) PrintElem() {
	fmt.Println("Level: ", sl.level)
	for i := sl.level; i >= 0; i-- {
		for cur := sl.head.next[i]; cur != nil; cur = cur.next[i] {
			fmt.Printf("(%v, %v) ", cur.key, cur.val)
		}
		fmt.Println()
	}
}

// Execute 执行各类操作的函数
func (sl *SkipList) Execute(op string, key SLKey, val interface{}) interface{} {
	switch op[0] {
	case 'I':
		sl.InsertKey(key, val)
		sl.PrintElem()
		return nil
	case 'S':
		return sl.SearchKey(key)
	case 'D':
		sl.DeleteKey(key)
		sl.PrintElem()
		return nil
	default:
	}
	return nil
}

2. Trie（动态&&静态）

package main

import "fmt"

/*
* 静态字典树
 */
const (
	maxNodeNum = 1010
)

// Node 节点信息
type Node struct {
	childs [26]*Node
	flag   bool // 结尾标记
}

var nds [maxNodeNum]Node
var curIdx int
var root *Node

func NewNode() *Node {
	newNode := &nds[curIdx]
	curIdx++
	return newNode
}

func InsertNode(nd *Node, s string) {
	i := 0
	slen := len(s)
	for i < slen {
		if nd.childs[s[i]-'a'] == nil {
			nd.childs[s[i]-'a'] = NewNode()
		}
		nd = nd.childs[s[i]-'a']
		i++
	}
	nd.flag = true
}

func QueryNode(nd *Node, s string) bool {
	i := 0
	slen := len(s)
	for i < slen {
		if nd.childs[s[i]-'a'] == nil {
			return false
		}
		nd = nd.childs[s[i]-'a']
		i++
	}
	return nd.flag
}

/*
func main() {
	root := NewNode()
	InsertNode(root, "abc")
	InsertNode(root, "acd")
	InsertNode(root, "bb")
	fmt.Println(QueryNode(root, "b"))
	fmt.Println(QueryNode(root, "abc"))
	fmt.Println(QueryNode(root, "abcd"))
}
*/

/*
* 动态字典树
 */

func NewDyncNode() *Node {
	return &Node{
		childs: [26]*Node{},
		flag:   false,
	}
}

func InsertDyncNode(nd *Node, s string) {
	i := 0
	slen := len(s)
	for i < slen {
		if nd.childs[s[i]-'a'] == nil {
			nd.childs[s[i]-'a'] = NewDyncNode()
		}
		nd = nd.childs[s[i]-'a']
		i++
	}
	nd.flag = true
}

func QueryDyncNode(nd *Node, s string) bool {
	i := 0
	slen := len(s)
	for i < slen {
		if nd.childs[s[i]-'a'] == nil {
			return false
		}
		nd = nd.childs[s[i]-'a']
		i++
	}
	return nd.flag
}

func main() {
	root := NewDyncNode()
	InsertDyncNode(root, "abc")
	InsertDyncNode(root, "acd")
	InsertDyncNode(root, "bb")
	fmt.Println(QueryDyncNode(root, "b"))
	fmt.Println(QueryDyncNode(root, "abc"))
	fmt.Println(QueryDyncNode(root, "abcd"))
	fmt.Println(QueryDyncNode(root, "a"))
	fmt.Println(QueryDyncNode(root, "acd"))
}

3. Radix Tree（Trie 变种）

package main

import (
	"fmt"
)

/*
* 简单Radix-Tree
* 只需要简单的插入、查询、删除、子树遍历操作即可
 */

type Entry struct {
	key string
	val string
}

type Node struct {
	leaf   *Entry
	prefix string
	childs map[byte]*Node
}

func NewNode(prefix string, entry *Entry) *Node {
	return &Node{
		leaf:   entry,
		prefix: prefix,
		childs: make(map[byte]*Node),
	}
}

func findCommonPrefix(s1, s2 string) int {
	idx := 0
	for ; idx < len(s1) && idx < len(s2) && s1[idx] == s2[idx]; idx++ {
	}
	return idx
}

// InsertNode 根据s插入节点entry
func InsertNode(root *Node, s string, entry *Entry) {

	var p, nd *Node
	var idx int
	nd = root
	for {

		if _, ok := nd.childs[s[0]]; !ok {
			nd.childs[s[0]] = NewNode(s, entry)
			return
		}

		p = nd
		nd = nd.childs[s[0]]
		idx = findCommonPrefix(nd.prefix, s)

		// 1. 最长公共前缀长度小于任何一个字符串，则节点需要分裂
		if idx < len(nd.prefix) && idx < len(s) {

			nd1 := NewNode(s[:idx], nil)
			p.childs[s[0]] = nd1

			nd1.childs[nd.prefix[idx]] = nd
			nd.prefix = nd.prefix[idx:]

			nd2 := NewNode(s[idx:], entry)
			nd1.childs[s[idx]] = nd2
			return
		}

		// 2.
		if idx >= len(s) && idx >= len(nd.prefix) {
			nd.leaf = entry
			return
		}

		// 3.
		if idx >= len(s) {
			nd1 := NewNode(s[:idx], entry)
			p.childs[s[0]] = nd1
			nd1.childs[nd.prefix[idx]] = nd
			nd.prefix = nd.prefix[idx:]
			return
		}

		// 4.
		s = s[idx:]
	}
}

// GetNode 根据s查找对应的节点
func GetNode(root *Node, s string) (bool, *Entry) {

	for {

		if _, ok := root.childs[s[0]]; !ok {
			return false, nil
		}

		root = root.childs[s[0]]
		idx := findCommonPrefix(s, root.prefix)

		if idx < len(root.prefix) {
			return false, nil
		}

		if idx == len(root.prefix) && idx == len(s) {
			return root.leaf != nil, root.leaf
		}

		s = s[idx:]
	}
	return false, nil
}

// DeleteNode 返回该节点是否存在以及该旧节点的数据
// 思路: 先定位节点然后, 再决定是否删除该节点, 删除过程中牵扯到节点合并
// 1. 直接干掉该节点：该节点无任何孩子节点
// 2. 合并该节点和其孩子节点：该节点不是根节点且该节点只有一个孩子节点
// 3. 合并父亲节点和该节点: 父节点不为根节点，父节点不是叶子节点，父节点的孩子节点只有一个，父节点不为空
func DeleteNode(root *Node, s string) {

	var isExist bool
	var p, nd *Node
	nd = root
	for {

		if _, ok := nd.childs[s[0]]; !ok {
			isExist = false
			break
		}

		p = nd
		nd = nd.childs[s[0]]
		idx := findCommonPrefix(s, nd.prefix)

		if idx < len(nd.prefix) {
			isExist = false
			break
		}

		if idx == len(nd.prefix) && idx == len(s) {
			if nd.leaf != nil {
				isExist = true
			} else {
				isExist = false
			}
			break
		}

		s = s[idx:]
	}

	if !isExist {
		return
	}

	// 置空该节点的值
	nd.leaf = nil

	// 1.
	if p != nil && len(nd.childs) == 0 {
		delete(p.childs, nd.prefix[0])
	}

	// 2.
	if nd != root && len(nd.childs) == 1 {
		for lb := range nd.childs {
			nd.leaf = nd.childs[lb].leaf
			nd.prefix = nd.prefix + nd.childs[lb].prefix
			nd.childs = nd.childs[lb].childs
			break
		}
	}

	// 3.
	if p != nil && p != root && len(p.childs) == 1 && p.leaf == nil {
		for lb := range p.childs {
			p.leaf = p.childs[lb].leaf
			p.prefix = p.prefix + p.childs[lb].prefix
			p.childs = p.childs[lb].childs
			break
		}
	}
	return
}

/*
func main() {
	root := NewNode("", nil)
	InsertNode(root, "abc", &Entry{key: "abc", val: "abc"})
	//ok1, nd1 := GetNode(root, "abc")
	//fmt.Println(ok1, nd1.key, nd1.val)
	//ok2, _ := GetNode(root, "ab")
	//fmt.Println(ok2)
	InsertNode(root, "abe", &Entry{key: "abe", val: "abe"})
	//_, ok := root.childs['a']
	//fmt.Println(ok)
	ok3, nd3 := GetNode(root, "abe")
	fmt.Println(ok3, nd3.key, nd3.val)

	InsertNode(root, "acd", &Entry{key: "acd", val: "acd"})
	ok4, nd4 := GetNode(root, "acd")
	fmt.Println(ok4, nd4.key, nd4.val)
	ok5, nd5 := GetNode(root, "abc")
	fmt.Println(ok5, nd5.key, nd5.val)
	ok6, nd6 := GetNode(root, "abe")
	fmt.Println(ok6, nd6.key, nd6.val)

	InsertNode(root, "b", &Entry{key: "b", val: "b"})
	ok7, nd7 := GetNode(root, "b")
	fmt.Println(ok7, nd7.key, nd7.val)
	//ok4, nd4 := GetNode(root, "abc")
	//fmt.Println(ok4, nd4.key, nd4.val)
	DeleteNode(root, "b")
	ok8, nd8 := GetNode(root, "b")
	fmt.Println(ok8, nd8)
}
*/

func main() {
	root := NewNode("", nil)
	InsertNode(root, "a", &Entry{key: "a", val: "a"})
	InsertNode(root, "ab", &Entry{key: "ab", val: "ab"})
	DeleteNode(root, "ab")
	ok1, nd1 := GetNode(root, "ab")
	fmt.Println(ok1, nd1)
	ok2, nd2 := GetNode(root, "a")
	fmt.Println(ok2, nd2)
	InsertNode(root, "abcd", &Entry{key: "abcd", val: "abcd"})
	InsertNode(root, "abc", &Entry{key: "abc", val: "abc"})
	DeleteNode(root, "abc")
	ok3, nd3 := GetNode(root, "abcd")
	fmt.Println(ok3, nd3)
	InsertNode(root, "abcde", &Entry{key: "abcde", val: "abcde"})
	InsertNode(root, "abcdf", &Entry{key: "abcdf", val: "abcdf"})
	DeleteNode(root, "abcd")
	ok4, nd4 := GetNode(root, "a")
	fmt.Println(ok4, nd4)
	ok5, nd5 := GetNode(root, "abcde")
	fmt.Println(ok5, nd5)
	ok6, nd6 := GetNode(root, "abcdf")
	fmt.Println(ok6, nd6)
	InsertNode(root, "abc", &Entry{key: "abc", val: "abc"})
	InsertNode(root, "abcd", &Entry{key: "abcd", val: "abcd"})
	DeleteNode(root, "abc")
	ok7, nd7 := GetNode(root, "abcde")
	fmt.Println(ok7, nd7)
	DeleteNode(root, "abcd")
	ok8, nd8 := GetNode(root, "a")
	fmt.Println(ok8, nd8)
}

/*
func main() {
	a := map[int]int{
		1: 1,
		2: 2,
	}
	for lb := range a {
		fmt.Println(lb, a[lb])
		a = map[int]int{
			3: 3,
			4: 4,
		}
		break
	}
	fmt.Println(a)
}
*/

4. 并查集（带权）

package main

import "fmt"

const (
	maxNum = 1010
)

var parent [maxNum]int

func initSet(n int) {
	for i := 0; i <= n; i++ {
		parent[i] = -1
	}
}

// C++ 一行代码 : return parent[x] < 0 ? x : parent[x] = find_set(parent[x])
func findSet(x int) int {
	if parent[x] < 0 {
		return x
	}
	parent[x] = findSet(parent[x])
	return parent[x]
}

func unionSet(x, y int) {
	p1 := findSet(x)
	p2 := findSet(y)
	if p1 == p2 {
		return
	}
	if parent[p1] < parent[p2] {
		parent[p1] += parent[p2]
		parent[p2] = p1
	} else {
		parent[p2] += parent[p1]
		parent[p1] = p2
	}
	return
}

// parent[i] < 0 : 该不相交集合的根节点

func main() {
	initSet(5)
	unionSet(0, 2)
	unionSet(1, 3)
	unionSet(2, 0)
	unionSet(3, 4)
	unionSet(4, 1)
	for i := 0; i < 5; i++ {
		if parent[i] < 0 {
			fmt.Println(i, -parent[i])
		}
	}
}

/*
* 带权并查集: 每个节点均存在权重，并查集维护集合中每个节点和根节点的相对关系，在集合合并的时候维护好集合
* 根节点之间的相对关系，且在查找过程的路径压缩时，需要更新和新的父节点的相对权重关系
* var parent [maxNode]int
* var val [maxNode]int
 */