146 lines
2.3 KiB
Go
146 lines
2.3 KiB
Go
package avl
|
|
|
|
import (
|
|
"github.com/emirpasic/gods/stacks/arraystack"
|
|
)
|
|
|
|
type Iterator struct {
|
|
op *AVL
|
|
|
|
dir int
|
|
up *Node
|
|
cur *Node
|
|
tstack *arraystack.Stack
|
|
// curnext *Node
|
|
}
|
|
|
|
func initIterator(avltree *AVL) *Iterator {
|
|
iter := &Iterator{op: avltree, tstack: arraystack.New()}
|
|
iter.up = avltree.root
|
|
return iter
|
|
}
|
|
|
|
func (iter *Iterator) Value() interface{} {
|
|
return iter.cur.value
|
|
}
|
|
|
|
func (iter *Iterator) Next() (result bool) {
|
|
|
|
if iter.dir < 1 {
|
|
if iter.dir == -1 {
|
|
if iter.cur != nil {
|
|
iter.tstack.Clear()
|
|
iter.curPushNextStack(iter.cur)
|
|
|
|
iter.up = iter.cur
|
|
for iter.up != nil {
|
|
if iter.up.child == 1 {
|
|
iter.up = iter.up.parent
|
|
break
|
|
} else {
|
|
iter.up = iter.up.parent
|
|
}
|
|
}
|
|
|
|
}
|
|
}
|
|
iter.dir = 1
|
|
}
|
|
|
|
if iter.tstack.Size() == 0 {
|
|
if iter.up != nil {
|
|
iter.tstack.Push(iter.up)
|
|
for iter.up != nil {
|
|
if iter.up.child == 1 {
|
|
iter.up = iter.up.parent
|
|
break
|
|
} else {
|
|
iter.up = iter.up.parent
|
|
}
|
|
}
|
|
} else {
|
|
return false
|
|
}
|
|
}
|
|
|
|
if v, ok := iter.tstack.Pop(); ok {
|
|
iter.cur = v.(*Node)
|
|
iter.curPushNextStack(iter.cur)
|
|
return true
|
|
}
|
|
|
|
return false
|
|
}
|
|
|
|
func (iter *Iterator) curPushNextStack(cur *Node) {
|
|
next := cur.children[0]
|
|
|
|
if next != nil {
|
|
iter.tstack.Push(next)
|
|
for next.children[1] != nil {
|
|
next = next.children[1]
|
|
iter.tstack.Push(next)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (iter *Iterator) curPushPrevStack(cur *Node) {
|
|
prev := cur.children[1]
|
|
|
|
if prev != nil {
|
|
iter.tstack.Push(prev)
|
|
for prev.children[0] != nil {
|
|
prev = prev.children[0]
|
|
iter.tstack.Push(prev)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (iter *Iterator) Prev() (result bool) {
|
|
|
|
if iter.dir > -1 {
|
|
|
|
if iter.dir == 1 {
|
|
if iter.cur != nil {
|
|
iter.tstack.Clear()
|
|
iter.curPushPrevStack(iter.cur)
|
|
|
|
iter.up = iter.cur
|
|
for iter.up != nil {
|
|
if iter.up.child == 0 {
|
|
iter.up = iter.up.parent
|
|
break
|
|
} else {
|
|
iter.up = iter.up.parent
|
|
}
|
|
}
|
|
}
|
|
}
|
|
iter.dir = -1
|
|
}
|
|
|
|
if iter.tstack.Size() == 0 {
|
|
if iter.up != nil {
|
|
iter.tstack.Push(iter.up)
|
|
for iter.up != nil {
|
|
if iter.up.child == 0 {
|
|
iter.up = iter.up.parent
|
|
break
|
|
} else {
|
|
iter.up = iter.up.parent
|
|
}
|
|
}
|
|
} else {
|
|
return false
|
|
}
|
|
}
|
|
|
|
if v, ok := iter.tstack.Pop(); ok {
|
|
iter.cur = v.(*Node)
|
|
iter.curPushPrevStack(iter.cur)
|
|
return true
|
|
}
|
|
|
|
return false
|
|
}
|