SizeBalancedTreeMap

package SortedMap

import (
	"fmt"
	"strings"
	"testing"
)

type KEY interface {
	CompareTo(v KEY) int
}

type VALUE interface{}

type SBTNode struct {
	key   KEY
	value VALUE
	l    *SBTNode
	r    *SBTNode
	size int   // 不同的key的数量
}

func NewSBTNode(key KEY, value VALUE) *SBTNode {
	return &SBTNode{
		key:   key,
		value: value,
		size:  1,
	}
}

type SizeBalancedTreeMap struct {
	root *SBTNode
}

func NewSizeBalancedTreeMap() *SizeBalancedTreeMap {
	return &SizeBalancedTreeMap{}
}

func (SBTMap *SizeBalancedTreeMap) rightRotate(cur *SBTNode) *SBTNode {
	leftNode  := cur.l
	cur.l = leftNode.r
	leftNode.r = cur
	leftNode.size = cur.size
	cur.size = 0
	if cur.l != nil {
		cur.size += cur.l.size
	}
	if cur.r != nil {
		cur.size += cur.r.size
	}
	cur.size += 1
	return leftNode
}

func (SBTMap *SizeBalancedTreeMap) leftRotate(cur *SBTNode) *SBTNode {
	rightNode := cur.r
	cur.r = rightNode.l
	rightNode.l = cur
	rightNode.size = cur.size
	cur.size = 0
	if cur.l != nil {
		cur.size += cur.l.size
	}
	if cur.r != nil {
		cur.size += cur.r.size
	}
	cur.size += 1
	return rightNode
}


func (SBTMap *SizeBalancedTreeMap) maintain(cur *SBTNode) *SBTNode {
	if cur == nil {
		return nil
	}
	leftSize := 0
	if cur.l != nil {
		leftSize = cur.l.size
	}
	leftLeftSize := 0
	if cur.l != nil && cur.l.l != nil {
		leftLeftSize = cur.l.l.size
	}
	leftRightSize := 0
	if cur.l != nil && cur.l.r != nil {
		leftRightSize = cur.l.r.size
	}
	rightSize := 0
	if cur.r != nil {
		rightSize = cur.r.size
	}
	rightLeftSize := 0
	if cur.r != nil && cur.r.l != nil {
		rightLeftSize = cur.r.l.size
	}
	rightRightSize := 0
	if cur.r != nil && cur.r.r != nil {
		rightRightSize = cur.r.r.size
	}

	if leftLeftSize > rightSize {
		cur = SBTMap.rightRotate(cur)
		cur.r =  SBTMap.rightRotate(cur.r)
		cur  =  SBTMap.maintain(cur)
	} else if  leftRightSize > rightSize {
		cur.l = SBTMap.leftRotate(cur.l)
		cur = SBTMap.rightRotate(cur)
		cur.l = SBTMap.maintain(cur.l)
		cur.r = SBTMap.maintain(cur.r)
		cur = SBTMap.maintain(cur)
	} else if  rightRightSize > leftSize  {
		cur = SBTMap.leftRotate(cur)
		cur.l = SBTMap.maintain(cur.l)
		cur = SBTMap.maintain(cur)
	} else if  rightLeftSize > leftSize  {
		cur.r = SBTMap.rightRotate(cur.r)
		cur = SBTMap.leftRotate(cur)
		cur.l = SBTMap.maintain(cur.l)
		cur.r = SBTMap.maintain(cur.r)
		cur = SBTMap.maintain(cur)
	}
	return cur
}


func (SBTMap *SizeBalancedTreeMap) FindLastIndex(key KEY) *SBTNode {
	pre := SBTMap.root
	cur := SBTMap.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 (SBTMap *SizeBalancedTreeMap) FindLastNoSmallIndex(key KEY) *SBTNode {
	ans, cur := (*SBTNode)(nil), SBTMap.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 (SBTMap *SizeBalancedTreeMap) FindLastNoBigIndex(key KEY) *SBTNode {
	ans, cur := (*SBTNode)(nil), SBTMap.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
}

// 现在，以cur为头的树上，新增，加(key, value)这样的记录
// 加完之后，会对cur做检查，该调整调整
// 返回，调整完之后，整棵树的新头部

func (SBTMap *SizeBalancedTreeMap) Add(cur *SBTNode, key KEY, value VALUE) *SBTNode {
	if cur == nil {
		return NewSBTNode(key, value)
	} else {
		cur.size++
		if key.CompareTo(cur.key) < 0 {
			cur.l = SBTMap.Add(cur.l, key, value)
		} else {
			cur.r = SBTMap.Add(cur.r, key, value)
		}
		return SBTMap.maintain(cur)
	}
}

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

func (SBTMap *SizeBalancedTreeMap) Delete(cur *SBTNode, key KEY) *SBTNode {
	cur.size--
	if key.CompareTo(cur.key) > 0 {
		cur.r = SBTMap.Delete(cur.r, key)
	} else if key.CompareTo(cur.key) < 0 {
		cur.l = SBTMap.Delete(cur.l, key)
	} else { // 当前要删掉cur
		if cur.l == nil && cur.r == nil {
			// free cur memory -> C++
			cur = nil
		} else if cur.l == nil && cur.r != nil {
			// free cur memory -> C++
			cur = cur.r
		} else if cur.l != nil && cur.r == nil {
			// free cur memory -> C++
			cur = cur.l
		} else { // 有左有右
			pre := (*SBTNode)(nil)
			des := cur.r
			des.size--
			for des.l != nil {
				pre = des
				des = des.l
				des.size--
			}
			if pre != nil {
				pre.l = des.r
				des.r = cur.r
			}
			des.l = cur.l

			des.size = des.l.size
			if des.r != nil {
				des.size += des.r.size
			}
			des.size += 1
			// free cur memory -> C++
			cur = des
		}
	}
	// cur = maintain(cur);
	return cur
}


func (BSTMap *SizeBalancedTreeMap) GetIndex(cur *SBTNode, kth int) *SBTNode {
	N := 0
	if cur.l != nil {
		N = cur.l.size
	}

	if kth == N + 1 {
		return cur
	} else if kth <= N {
	    return BSTMap.GetIndex(cur.l, kth)
	}

	return BSTMap.GetIndex(cur.r, kth - N - 1)
}

func (BSTMap *SizeBalancedTreeMap)Size() int {
	return BSTMap.root.size
}

func (BSTMap *SizeBalancedTreeMap) Contains(key KEY) bool {
	if key == nil {
		return false
	}
	lastNode := BSTMap.FindLastIndex(key)
	return lastNode != nil && key.CompareTo(lastNode.key) == 0
}

func (BSTMap *SizeBalancedTreeMap) Put(key KEY, value VALUE) {
	if key == nil {
		return
	}
	lastNode := BSTMap.FindLastIndex(key)
	if lastNode != nil && key.CompareTo(lastNode.key) == 0 {
		lastNode.value = value
	} else {
		BSTMap.root = BSTMap.Add(BSTMap.root, key, value)
	}
}

func (BSTMap *SizeBalancedTreeMap) Remove(key KEY) {
	if key == nil {
		return
	}
	if BSTMap.Contains(key) {
		BSTMap.root = BSTMap.Delete(BSTMap.root, key)
	}
}

func (BSTMap *SizeBalancedTreeMap) GetIndexKey(index int) KEY { // todo err检测
	if index < 0 || index >= BSTMap.Size() {
		return nil
	}
	return BSTMap.GetIndex(BSTMap.root, index + 1).key
}


func (BSTMap *SizeBalancedTreeMap) GetIndexValue(index int) VALUE {
	if index < 0 || index >= BSTMap.Size() {
		return nil
	}
	return BSTMap.GetIndex(BSTMap.root, index + 1).value
}



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


func (BSTMap *SizeBalancedTreeMap) FirstKey() KEY {
	if BSTMap.root == nil {
		return nil
	}
	cur := BSTMap.root
	for cur.l != nil {
		cur = cur.l
	}
	return cur.key
}

func (BSTMap *SizeBalancedTreeMap) LastKey() KEY {
	if BSTMap.root == nil {
		return nil
	}
	cur := BSTMap.root
	for cur.r != nil {
		cur = cur.r
	}
	return cur.key
}



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




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



func (BSTMap *SizeBalancedTreeMap) printAll() {
	fmt.Println("Binary Tree:")
	BSTMap.printInOrder(BSTMap.root, 0, "H", 17)
	fmt.Println()
}
func (BSTMap *SizeBalancedTreeMap)printInOrder(head *SBTNode, height int, to string, length int) {
	if head == nil {
		return
	}
	BSTMap.printInOrder(head.r, height + 1, "v", length)
	val := to + fmt.Sprintf("(%t,%d)",head.key,head.value) + to
	lenM := len(val)
	lenL := (length - lenM) / 2
	lenR := length - lenM - lenL
	val = BSTMap.getSpace(lenL) + val + BSTMap.getSpace(lenR)
	fmt.Println(BSTMap.getSpace(height * length) + val)
	BSTMap.printInOrder(head.l, height + 1, "^", length)
}

func (BSTMap *SizeBalancedTreeMap)getSpace(num int) string {
	space := " "
	buf := strings.Builder{}
	for i := 0; i < num; i++ {
		buf.WriteString(space)
	}
	return buf.String()
}


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 TestSizeBalancedTree(t *testing.T)  {
	sbt := NewSizeBalancedTreeMap()

	s1 := &student{"张1",0}
	s2 := &student{"张2",1}
	s3 := &student{"张3",2}
	s4 := &student{"张4",3}
	s5 := &student{"张5",4}
	s6 := &student{"张6",5}

	sbt.Put(s1, 4)
	sbt.Put(s2, 3)
	sbt.Put(s3, 1)
	sbt.Put(s4, 2)
	sbt.Put(s5, 7)
	sbt.Put(s6, 6)

	fmt.Println(sbt.Get(s1))
	sbt.Put(s1,100)
	fmt.Println(sbt.Get(s1))

	for  i := 0; i < sbt.Size(); i++ {
		fmt.Println(sbt.GetIndexKey(i), sbt.GetIndexValue(i))
	}
	sbt.printAll()
	fmt.Println(sbt.FirstKey())
	fmt.Println(sbt.LastKey())
	fmt.Println(sbt.FloorKey(s6))
	fmt.Println(sbt.CeilingKey(s2))
	fmt.Println(sbt.FloorKey(s5))
	fmt.Println(sbt.CeilingKey(s4))
	fmt.Println(sbt.FloorKey(s1))
	fmt.Println(sbt.CeilingKey(s6))
	fmt.Println(sbt.FloorKey(s2))
	fmt.Println(sbt.CeilingKey(s3))
	sbt.Remove(s2)
	sbt.printAll()
	sbt.Remove(s5)
	sbt.printAll()
}

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