213

avl/avl.go

@@ -1,9 +1,8 @@
 package avl
 
 import (
+	"474420502.top/eson/structure/compare"
 	"github.com/davecgh/go-spew/spew"
-
-	"github.com/emirpasic/gods/utils"
 )
 
 type Node struct {
@@ -26,13 +25,13 @@ func (n *Node) String() string {
 }
 
 type Tree struct {
-	root       *Node
-	size       int
-	comparator utils.Comparator
+	root    *Node
+	size    int
+	compare compare.Compare
 }
 
-func New(comparator utils.Comparator) *Tree {
-	return &Tree{comparator: comparator}
+func New(compare compare.Compare) *Tree {
+	return &Tree{compare: compare}
 }
 
 func (avl *Tree) String() string {
@@ -118,18 +117,83 @@ func (avl *Tree) Remove(key interface{}) *Node {
 }
 
 // Values 返回先序遍历的值
-func (avl *Tree) Values() []interface{} {
+func (tree *Tree) Values() []interface{} {
 	mszie := 0
-	if avl.root != nil {
-		mszie = avl.size
+	if tree.root != nil {
+		mszie = tree.size
 	}
 	result := make([]interface{}, 0, mszie)
-	avl.Traversal(func(v interface{}) {
+	tree.Traversal(func(v interface{}) bool {
 		result = append(result, v)
-	}, DLR)
+		return true
+	}, LDR)
 	return result
 }
 
+func (tree *Tree) GetRange(k1, k2 interface{}) (result []interface{}) {
+	c := tree.compare(k2, k1)
+	switch c {
+	case 1:
+
+		var min, max *Node
+		resultmin := tree.getArountNode(k1)
+		resultmax := tree.getArountNode(k2)
+		for i := 1; i < 3 && min == nil; i++ {
+			min = resultmin[i]
+		}
+
+		for i := 1; i > -1 && max == nil; i-- {
+			max = resultmax[i]
+		}
+
+		if max == nil {
+			return []interface{}{}
+		}
+
+		result = make([]interface{}, 0, 16)
+
+		iter := NewIterator(min)
+		for iter.Prev() {
+			result = append(result, iter.Value())
+			if iter.cur == max {
+				break
+			}
+		}
+	case -1:
+
+		var min, max *Node
+		resultmin := tree.getArountNode(k2)
+		resultmax := tree.getArountNode(k1)
+		for i := 1; i < 3 && min == nil; i++ {
+			min = resultmin[i]
+		}
+		for i := 1; i > -1 && max == nil; i-- {
+			max = resultmax[i]
+		}
+
+		if min == nil {
+			return []interface{}{}
+		}
+
+		result = make([]interface{}, 0, 16)
+
+		iter := NewIterator(max)
+		for iter.Next() {
+			result = append(result, iter.Value())
+			if iter.cur == min {
+				break
+			}
+		}
+	case 0:
+		if n, ok := tree.GetNode(k1); ok {
+			return []interface{}{n.value}
+		}
+		return []interface{}{}
+	}
+
+	return
+}
+
 func (avl *Tree) Get(key interface{}) (interface{}, bool) {
 	n, ok := avl.GetNode(key)
 	if ok {
@@ -138,12 +202,8 @@ func (avl *Tree) Get(key interface{}) (interface{}, bool) {
 	return n, false
 }
 
-func (avl *Tree) GetRange(min, max interface{}) []interface{} {
-	return nil
-}
-
-func (avl *Tree) GetAround(key interface{}) (result [3]interface{}) {
-	an := avl.GetAroundNode(key)
+func (tree *Tree) GetAround(key interface{}) (result [3]interface{}) {
+	an := tree.getArountNode(key)
 	for i, n := range an {
 		if n != nil {
 			result[i] = n.value
@@ -152,41 +212,165 @@ func (avl *Tree) GetAround(key interface{}) (result [3]interface{}) {
 	return
 }
 
-func (avl *Tree) GetAroundNode(value interface{}) (result [3]*Node) {
-	if cur, ok := avl.GetNode(value); ok {
+func (tree *Tree) getArountNode(key interface{}) (result [3]*Node) {
+	var last *Node
+	var lastc int
 
-		var iter *Iterator
+	for n := tree.root; n != nil; {
+		last = n
+		c := tree.compare(key, n.value)
+		switch c {
+		case -1:
+			n = n.children[0]
+			lastc = c
+		case 1:
+			n = n.children[1]
+			lastc = c
+		case 0:
+			result[1] = n
+			n = nil
+		default:
+			panic("Get compare only is allowed in -1, 0, 1")
+		}
+	}
 
-		iter = NewIterator(cur)
-		iter.curPushPrevStack(iter.up)
-		iter.up = iter.getPrevUp(iter.up)
+	switch lastc {
+	case 1:
+		const il = 0
+		const ir = 1
+
+		if result[1] == nil {
+			result[0] = last
+
+			parent := last
+			for ; parent != nil && parent.parent != nil; parent = parent.parent {
+				child := getRelationship(parent)
+				if child == (-lastc+2)/2 { // child 与 compare 后左右的关系
+					result[2] = parent.parent
+					break
+				}
+			}
 
-		if v, ok := iter.tstack.Pop(); ok {
-			result[0] = v.(*Node)
-			// iter.curPushPrevStack(iter.cur)
 		} else {
-			result[0] = iter.up
+			l := result[1].children[il]
+			r := result[1].children[ir]
+
+			if l == nil {
+				result[0] = result[1].parent
+			} else {
+				for l.children[ir] != nil {
+					l = l.children[ir]
+				}
+				result[0] = l
+			}
+
+			if r == nil {
+
+				parent := result[1].parent
+				for ; parent != nil && parent.parent != nil; parent = parent.parent {
+					child := getRelationship(parent)
+					if child == (-lastc+2)/2 { // child 与 compare 后左右的关系
+						result[2] = parent.parent
+						break
+					}
+				}
+
+			} else {
+				for r.children[il] != nil {
+					r = r.children[il]
+				}
+				result[2] = r
+			}
 		}
 
-		iter = NewIterator(cur)
-		iter.curPushNextStack(iter.up)
-		iter.up = iter.getNextUp(iter.up)
+	case -1:
+
+		const il = 1
+		const ir = 0
+
+		if result[1] == nil {
+			result[2] = last
+
+			parent := last
+			for ; parent != nil && parent.parent != nil; parent = parent.parent {
+				child := getRelationship(parent)
+				if child == (-lastc+2)/2 { // child 与 compare 后左右的关系
+					result[0] = parent.parent
+					break
+				}
+			}
 
-		if v, ok := iter.tstack.Pop(); ok {
-			result[2] = v.(*Node)
-			// iter.curPushNextStack(iter.cur)
 		} else {
-			result[2] = iter.up
+
+			l := result[1].children[il]
+			r := result[1].children[ir]
+
+			if l == nil {
+				result[2] = result[1].parent
+			} else {
+				for l.children[ir] != nil {
+					l = l.children[ir]
+				}
+				result[2] = l
+			}
+
+			if r == nil {
+
+				parent := result[1].parent
+				for ; parent != nil && parent.parent != nil; parent = parent.parent {
+					child := getRelationship(parent)
+					if child == (-lastc+2)/2 { // child 与 compare 后左右的关系
+						result[0] = parent.parent
+						break
+					}
+				}
+
+			} else {
+
+				for r.children[il] != nil {
+					r = r.children[il]
+				}
+				result[0] = r
+
+			}
+		}
+	case 0:
+
+		const il = 0
+		const ir = 1
+
+		if result[1] == nil {
+			return
 		}
 
-		result[1] = cur
+		l := result[1].children[il]
+		r := result[1].children[ir]
+
+		if l == nil {
+			result[0] = nil
+		} else {
+			for l.children[ir] != nil {
+				l = l.children[ir]
+			}
+			result[0] = l
+		}
+
+		if r == nil {
+			result[2] = nil
+		} else {
+			for r.children[il] != nil {
+				r = r.children[il]
+			}
+			result[2] = r
+		}
 	}
 	return
 }
+
 func (avl *Tree) GetNode(value interface{}) (*Node, bool) {
 
 	for n := avl.root; n != nil; {
-		switch c := avl.comparator(value, n.value); c {
+		switch c := avl.compare(value, n.value); c {
 		case -1:
 			n = n.children[0]
 		case 1:
@@ -194,7 +378,7 @@ func (avl *Tree) GetNode(value interface{}) (*Node, bool) {
 		case 0:
 			return n, true
 		default:
-			panic("Get comparator only is allowed in -1, 0, 1")
+			panic("Get compare only is allowed in -1, 0, 1")
 		}
 	}
 	return nil, false
@@ -224,36 +408,37 @@ func (avl *Tree) Put(value interface{}) {
 		}
 
 		parent = cur
-		c := avl.comparator(value, cur.value)
+		c := avl.compare(value, cur.value)
 		child = (c + 2) / 2
 		cur = cur.children[child]
 	}
 }
 
-func (avl *Tree) debugString() string {
-	if avl.size == 0 {
-		return ""
-	}
-	str := "AVLTree\n"
-	outputfordebug(avl.root, "", true, &str)
-	return str
-}
-
 type TraversalMethod int
 
 const (
+	// L = left R = right D = Value(dest)
 	_ TraversalMethod = iota
-	//DLR 前序遍历
+	//DLR 先值 然后左递归 右递归 下面同理
 	DLR
-	//LDR 中序遍历
+	//LDR 先从左边有序访问到右边 从小到大
 	LDR
-	//LRD 后序遍历
+	// LRD 同理
 	LRD
+
+	// DRL 同理
+	DRL
+
+	// RDL 先从右边有序访问到左边 从大到小
+	RDL
+
+	// RLD 同理
+	RLD
 )
 
-// Traversal 遍历的方法
-func (avl *Tree) Traversal(every func(v interface{}), traversalMethod ...interface{}) {
-	if avl.root == nil {
+// Traversal 遍历的方法 默认是LDR 从小到大 compare 为 l < r
+func (tree *Tree) Traversal(every func(v interface{}) bool, traversalMethod ...interface{}) {
+	if tree.root == nil {
 		return
 	}
 
@@ -264,40 +449,114 @@ func (avl *Tree) Traversal(every func(v interface{}), traversalMethod ...interfa
 
 	switch method {
 	case DLR:
-		var traverasl func(cur *Node)
-		traverasl = func(cur *Node) {
+		var traverasl func(cur *Node) bool
+		traverasl = func(cur *Node) bool {
 			if cur == nil {
-				return
+				return true
 			}
-			traverasl(cur.children[0])
-			every(cur.value)
-			traverasl(cur.children[1])
-		}
-		traverasl(avl.root)
-	case LRD:
-		var traverasl func(cur *Node)
-		traverasl = func(cur *Node) {
-			if cur == nil {
-				return
+			if !every(cur.value) {
+				return false
 			}
-			traverasl(cur.children[1])
-			every(cur.value)
-			traverasl(cur.children[0])
+			if !traverasl(cur.children[0]) {
+				return false
+			}
+			if !traverasl(cur.children[1]) {
+				return false
+			}
+			return true
 		}
-		traverasl(avl.root)
+		traverasl(tree.root)
 	case LDR:
-		var traverasl func(cur *Node)
-		traverasl = func(cur *Node) {
+		var traverasl func(cur *Node) bool
+		traverasl = func(cur *Node) bool {
 			if cur == nil {
-				return
+				return true
 			}
-			every(cur.value)
-			traverasl(cur.children[0])
-			traverasl(cur.children[1])
+			if !traverasl(cur.children[0]) {
+				return false
+			}
+			if !every(cur.value) {
+				return false
+			}
+			if !traverasl(cur.children[1]) {
+				return false
+			}
+			return true
 		}
-		traverasl(avl.root)
+		traverasl(tree.root)
+	case LRD:
+		var traverasl func(cur *Node) bool
+		traverasl = func(cur *Node) bool {
+			if cur == nil {
+				return true
+			}
+			if !traverasl(cur.children[0]) {
+				return false
+			}
+			if !traverasl(cur.children[1]) {
+				return false
+			}
+			if !every(cur.value) {
+				return false
+			}
+			return true
+		}
+		traverasl(tree.root)
+	case DRL:
+		var traverasl func(cur *Node) bool
+		traverasl = func(cur *Node) bool {
+			if cur == nil {
+				return true
+			}
+			if !every(cur.value) {
+				return false
+			}
+			if !traverasl(cur.children[0]) {
+				return false
+			}
+			if !traverasl(cur.children[1]) {
+				return false
+			}
+			return true
+		}
+		traverasl(tree.root)
+	case RDL:
+		var traverasl func(cur *Node) bool
+		traverasl = func(cur *Node) bool {
+			if cur == nil {
+				return true
+			}
+			if !traverasl(cur.children[1]) {
+				return false
+			}
+			if !every(cur.value) {
+				return false
+			}
+			if !traverasl(cur.children[0]) {
+				return false
+			}
+			return true
+		}
+		traverasl(tree.root)
+	case RLD:
+		var traverasl func(cur *Node) bool
+		traverasl = func(cur *Node) bool {
+			if cur == nil {
+				return true
+			}
+			if !traverasl(cur.children[1]) {
+				return false
+			}
+			if !traverasl(cur.children[0]) {
+				return false
+			}
+			if !every(cur.value) {
+				return false
+			}
+			return true
+		}
+		traverasl(tree.root)
 	}
-	return
 }
 
 func (avl *Tree) lrrotate(cur *Node) {
@@ -563,23 +822,14 @@ func getHeight(cur *Node) int {
 	return cur.height
 }
 
-func abs(n int) int {
-	y := n >> 31
-	return (n ^ y) - y
-}
-
 func (avl *Tree) fixRemoveHeight(cur *Node) {
 	for {
 
 		lefth, rigthh, lrmax := getMaxAndChildrenHeight(cur)
 
 		// 判断当前节点是否有变化, 如果没变化的时候, 不需要往上修复
-		isBreak := false
-		if cur.height == lrmax+1 {
-			isBreak = true
-		} else {
-			cur.height = lrmax + 1
-		}
+		curheight := lrmax + 1
+		cur.height = curheight
 
 		// 计算高度的差值 绝对值大于2的时候需要旋转
 		diff := lefth - rigthh
@@ -598,7 +848,7 @@ func (avl *Tree) fixRemoveHeight(cur *Node) {
 				avl.rrotate(cur)
 			}
 		} else {
-			if isBreak {
+			if cur.height == curheight {
 				return
 			}
 		}
@@ -722,3 +972,12 @@ func outputfordebug(node *Node, prefix string, isTail bool, str *string) {
 		outputfordebug(node.children[0], newPrefix, true, str)
 	}
 }
+
+func (avl *Tree) debugString() string {
+	if avl.size == 0 {
+		return ""
+	}
+	str := "AVLTree\n"
+	outputfordebug(avl.root, "", true, &str)
+	return str
+}

avl/avl_test.go

@@ -95,6 +95,7 @@ func TestIterator(t *testing.T) {
 	}
 
 	l = []int{21, 20, 15, 14, 7, 7, 6, 5, 4, 3, 2, 1}
+	t.Error(iter.cur, iter.tstack)
 	for i := 0; iter.Next(); i++ { // cur is 30 next is 21
 		if iter.Value().(int) != l[i] {
 			t.Error(iter.Value())

avl/iterator.go

@@ -1,12 +1,12 @@
 package avl
 
 import (
+	"log"
+
 	"474420502.top/eson/structure/lastack"
 )
 
 type Iterator struct {
-	op *Tree
-
 	dir    int
 	up     *Node
 	cur    *Node
@@ -15,7 +15,7 @@ type Iterator struct {
 }
 
 func initIterator(avltree *Tree) *Iterator {
-	iter := &Iterator{op: avltree, tstack: lastack.New()}
+	iter := &Iterator{tstack: lastack.New()}
 	iter.up = avltree.root
 	return iter
 }
@@ -48,6 +48,36 @@ func (iter *Iterator) Right() bool {
 	return false
 }
 
+func GetPrev(cur *Node, idx int) *Node {
+
+	iter := NewIterator(cur)
+	iter.curPushPrevStack(iter.up)
+	iter.up = iter.getPrevUp(iter.up)
+
+	for i := 0; i < idx; i++ {
+
+		if iter.tstack.Size() == 0 {
+			if iter.up == nil {
+				return nil
+			}
+			iter.tstack.Push(iter.up)
+			iter.up = iter.getPrevUp(iter.up)
+		}
+
+		if v, ok := iter.tstack.Pop(); ok {
+			iter.cur = v.(*Node)
+			if i == idx-1 {
+				return iter.cur
+			}
+			iter.curPushPrevStack(iter.cur)
+		} else {
+			return nil
+		}
+	}
+
+	return cur
+}
+
 func (iter *Iterator) Prev() (result bool) {
 
 	if iter.dir > -1 {
@@ -75,6 +105,35 @@ func (iter *Iterator) Prev() (result bool) {
 
 	return false
 }
+func GetNext(cur *Node, idx int) *Node {
+
+	iter := NewIterator(cur)
+	iter.curPushNextStack(iter.up)
+	iter.up = iter.getNextUp(iter.up)
+
+	for i := 0; i < idx; i++ {
+
+		if iter.tstack.Size() == 0 {
+			if iter.up == nil {
+				return nil
+			}
+			iter.tstack.Push(iter.up)
+			iter.up = iter.getNextUp(iter.up)
+		}
+
+		if v, ok := iter.tstack.Pop(); ok {
+			iter.cur = v.(*Node)
+			if i == idx-1 {
+				return iter.cur
+			}
+			iter.curPushNextStack(iter.cur)
+		} else {
+			return nil
+		}
+	}
+
+	return cur
+}
 
 func (iter *Iterator) Next() (result bool) {
 
@@ -127,6 +186,7 @@ func (iter *Iterator) curPushNextStack(cur *Node) {
 	next := cur.children[0] // 当前的左然后向右找, 找到最大, 就是最接近cur 并且小于cur的值
 
 	if next != nil {
+		log.Println(iter.tstack)
 		iter.tstack.Push(next)
 		for next.children[1] != nil {
 			next = next.children[1]