AVLTreeMap

package SortedMap

import "testing"

type KEY interface {
	CompareTo(v KEY) int
}

type VALUE interface{}

type AVLNode struct {
	key   KEY
	value VALUE
	l *AVLNode
	r *AVLNode
	h int
}

func NewAVLNode(key KEY, value VALUE) *AVLNode {
	return &AVLNode{
		key:   key,
		value: value,
		h:     1,
	}
}


type AVLTreeMap struct {
	Root *AVLNode
	size int
}

func NewAVLTreeMap() *AVLTreeMap {
	return &AVLTreeMap{}
}



func (AvlTMap *AVLTreeMap) rightRotate(cur *AVLNode) *AVLNode {
	left := cur.l
	cur.l = left.r
	left.r = cur


	lh, rh := 0, 0
	if cur.l != nil {
		lh = cur.l.h
	}
	if cur.r != nil {
		rh = cur.r.h
	}

	cur.h = Max(lh,rh) + 1

	lh, rh = 0, 0

	if left.l != nil {
		lh = left.l.h
	}
	if left.r != nil {
		rh = left.r.h
	}
	left.h = Max(lh,rh) + 1
	return left
}


func (AvlTMap *AVLTreeMap) leftRotate(cur *AVLNode) *AVLNode {
	right := cur.r
	cur.r = right.l
	right.l = cur


	lh, rh := 0, 0
	if cur.l != nil {
		lh = cur.l.h
	}
	if cur.r != nil {
		rh = cur.r.h
	}

	cur.h = Max(lh,rh) + 1

	lh, rh = 0, 0

	if right.l != nil {
		lh = right.l.h
	}
	if right.r != nil {
		rh = right.r.h
	}
	right.h = Max(lh,rh) + 1
	return right
}


func Max(a, b int) int {
	if a < b {
		return b
	}
	return a
}

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

func Abs(a int) int {
	if a < 0 {
		return  - a
	}
	return a
}




func (AvlTMap *AVLTreeMap)maintain(cur *AVLNode) *AVLNode {
	if cur == nil {
		return nil
	}

	leftHeight, rightHeight := 0, 0
	if cur.l != nil {
		leftHeight = cur.l.h
	}
	if cur.r != nil {
		rightHeight = cur.r.h
	}

	if Abs(leftHeight - rightHeight) > 1 {
		if leftHeight > rightHeight {
			leftLeftHeight, leftRightHeight := 0, 0
			if cur.l != nil && cur.l.l != nil {
				leftLeftHeight = cur.l.l.h
			}
			if cur.l != nil && cur.l.r != nil {
				leftRightHeight = cur.l.r.h
			}
			if leftLeftHeight >= leftRightHeight {
				cur = AvlTMap.rightRotate(cur)
			}else {
				cur.l = AvlTMap.leftRotate(cur.l)
				cur   = AvlTMap.rightRotate(cur)
			}
		}else {
			rightLeftHeight, rightRightHeight := 0, 0
			if cur.r != nil && cur.r.l != nil {
				rightLeftHeight = cur.r.l.h
			}
			if cur.r != nil && cur.r.r != nil {
				rightRightHeight = cur.r.r.h
			}
			if rightRightHeight >= rightLeftHeight {
				cur = AvlTMap.leftRotate(cur)
			}else {
				cur.r = AvlTMap.rightRotate(cur.r)
				cur   = AvlTMap.leftRotate(cur)
			}
		}
	}
	return cur
}

func (AvlTMap *AVLTreeMap)FindLastIndex(key KEY) *AVLNode {
	pre, cur := AvlTMap.Root, AvlTMap.Root
	for cur != nil {
		pre = cur
		if key.CompareTo(cur.key) == 0 {
			break
		}else if key.CompareTo(cur.key) < 0 {
			cur = cur.l
		}else {
			cur = cur.r
		}
	}
	return pre
}

func (AvlTMap *AVLTreeMap) FindLastNoSmallIndex(key KEY) *AVLNode {
	ans, cur := (*AVLNode)(nil), AvlTMap.Root
	for cur != nil {
		if key.CompareTo(cur.key) == 0 {
			ans = cur
			break
		}else if key.CompareTo(cur.key) < 0 {
			ans = cur
			cur = cur.l
		}else {
			cur = cur.r
		}
	}
	return ans
}

func (AvlTMap *AVLTreeMap) FindLastNoBigIndex(key KEY) *AVLNode {
	ans, cur := (*AVLNode)(nil), AvlTMap.Root
	for cur != nil {
		if key.CompareTo(cur.key) == 0 {
			ans = cur
			break
		}else if key.CompareTo(cur.key) < 0 {
			cur = cur.l
		}else {
			ans = cur
			cur = cur.r
		}
	}
	return ans
}


func (AvlTMap *AVLTreeMap) Add(cur *AVLNode,key KEY, value VALUE) *AVLNode {
	if cur == nil {
		return NewAVLNode(key,value)
	}else {
		if key.CompareTo(cur.key) < 0 {
			cur.l = AvlTMap.Add(cur.l,key,value)
		}else {
			cur.r = AvlTMap.Add(cur.r,key,value)
		}
		lh, rh := 0, 0
		if cur.l != nil {
			lh = cur.l.h
		}
		if cur.r != nil {
			rh = cur.r.h
		}

		cur.h = Max(lh,rh) + 1
		return AvlTMap.maintain(cur)
	}
}

// 在cur这棵树上，删掉key所代表的节点
// 返回cur这棵树的新头部

func (AvlTMap *AVLTreeMap) Delete(cur *AVLNode, key KEY) *AVLNode {
	if key.CompareTo(cur.key) > 0 {
		cur.r = AvlTMap.Delete(cur.r, key)
	} else if key.CompareTo(cur.key) < 0 {
		cur.l = AvlTMap.Delete(cur.l, key)
	} else {
		if cur.l == nil && cur.r == nil {
			cur = nil
		} else if cur.l == nil && cur.r != nil {
			cur = cur.r
		} else if cur.l != nil && cur.r == nil {
			cur = cur.l
		} else {
			 des := cur.r
			for des.l != nil {
				des = des.l
			}
			cur.r = AvlTMap.Delete(cur.r, des.key)
			des.l = cur.l
			des.r = cur.r
			cur = des
		}
	}
	if cur != nil {
		lh, rh := 0, 0
		if cur.l != nil {
			lh = cur.l.h
		}
		if cur.r != nil {
			rh = cur.r.h
		}
		cur.h = Max(lh,rh) + 1
	}
	return AvlTMap.maintain(cur)
}



func (AvlTMap *AVLTreeMap) Size() int {
	return AvlTMap.size
}

func (AvlTMap *AVLTreeMap) ContainsKey(key KEY) bool {
	if key == nil {
		return false
	}
	lastNode := AvlTMap.FindLastIndex(key)
	return lastNode != nil && key.CompareTo(lastNode.key) == 0
}


func (AvlTMap *AVLTreeMap) Put(key KEY, value VALUE) {
	if key == nil {
		return
	}
	lastNode := AvlTMap.FindLastNoSmallIndex(key)
	if lastNode != nil && key.CompareTo(lastNode.key) == 0 {
		lastNode.value = value
	}else {
		AvlTMap.size++
		AvlTMap.Root = AvlTMap.Add(AvlTMap.Root,key,value)
	}
}


func (AvlTMap *AVLTreeMap) Remove(key KEY) {
	if key == nil {
		return
	}
	if AvlTMap.ContainsKey(key) {
		AvlTMap.size--
		AvlTMap.Root = AvlTMap.Delete(AvlTMap.Root, key)
	}
}

func (AvlTMap *AVLTreeMap) Get(key KEY) VALUE {
	if key == nil {
		return nil
	}
	lastNode := AvlTMap.FindLastIndex(key)
	if lastNode != nil && key.CompareTo(lastNode.key) == 0 {
		return lastNode.value
	}
	return nil
}


func (AvlTMap *AVLTreeMap) FirstKey() KEY {
	if AvlTMap.Root == nil {
		return nil
	}
	cur := AvlTMap.Root
	for cur.l != nil {
		cur = cur.l
	}
	return cur.key
}

func (AvlTMap *AVLTreeMap) LastKey() KEY {
	if AvlTMap.Root == nil {
		return nil
	}
	cur := AvlTMap.Root
	for cur.r != nil {
		cur = cur.r
	}
	return cur.key
}

func (AvlTMap *AVLTreeMap) FloorKey(key KEY) KEY {
	if key == nil {
		return nil
	}
	lastNoBigNode := AvlTMap.FindLastNoBigIndex(key)
	if lastNoBigNode == nil {
		return nil
	}
	return lastNoBigNode.key
}

func (AvlTMap *AVLTreeMap) CeilingKey(key KEY) KEY {
	if key == nil {
		return nil
	}
	lastNoSmallNode := AvlTMap.FindLastNoSmallIndex(key)
	if lastNoSmallNode == nil {
		return nil
	}
    return lastNoSmallNode.key
}



type student struct {
	Name string
	Age  int
}

func (s *student)CompareTo(val KEY) int {
	if s.Name > val.(*student).Name {
		return 1
	}else if s.Name < val.(*student).Name {
		return -1
	}
	return 0
}

func TestAVLTreeMAP(t *testing.T)  {
	s1 := &student{"张1",0}
	s2 := &student{"张2",1}
	s3 := &student{"张3",2}
	s4 := &student{"张4",3}
	s5 := &student{"张5",4}
	s6 := &student{"张6",5}

	treeMap := NewAVLTreeMap()
	treeMap.Put(s1,"1")
	treeMap.Put(s2,"2")
	treeMap.Put(s3,"3")
	treeMap.Put(s4,"4")
	treeMap.Put(s5,"5")
	treeMap.Put(s6,"6")

	t.Log(treeMap.Get(s6))
	t.Log(treeMap.Get(s1))
	t.Log(treeMap.Get(s2))
	t.Log(treeMap.Get(s3))
	t.Log(treeMap.Get(s4))
	t.Log(treeMap.Get(s5))
	treeMap.Remove(s2)
	t.Log(treeMap.Get(s6))
	t.Log(treeMap.Get(s1))
	t.Log(treeMap.Get(s2))
	t.Log(treeMap.Get(s3))
	t.Log(treeMap.Get(s4))
	t.Log(treeMap.Get(s5))
	t.Log(treeMap.Size())
	t.Log(treeMap.CeilingKey(s2))
}

注：本人不是算法的原作者，我只是用Go重写了Java的代码，参照该地址，版权归原作者所有！