package pqueue import ( "474420502.top/eson/structure/lastack" ) type vbtIterator struct { dir int up *tNode cur *tNode tstack *lastack.Stack // curnext *tNode } func initIterator(avltree *vbTree) *vbtIterator { iter := &vbtIterator{tstack: lastack.New()} iter.up = avltree.root return iter } func NewIterator(n *tNode) *vbtIterator { iter := &vbtIterator{tstack: lastack.New()} iter.up = n return iter } func NewIteratorWithCap(n *tNode, cap int) *vbtIterator { iter := &vbtIterator{tstack: lastack.NewWithCap(cap)} iter.up = n return iter } func (iter *vbtIterator) SetNode(n *tNode) { iter.up = n iter.dir = 0 iter.tstack.Clear() } func (iter *vbtIterator) Value() interface{} { return iter.cur.value } func (iter *vbtIterator) Left() bool { if iter.cur.children[0] != nil { iter.dir = 0 iter.cur = iter.cur.children[0] return true } return false } func (iter *vbtIterator) Right() bool { if iter.cur.children[1] != nil { iter.dir = 0 iter.cur = iter.cur.children[1] return true } return false } func (iter *vbtIterator) GetNext(cur *tNode, idx int) *tNode { // iter := NewIterator(cur) iter.SetNode(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.(*tNode) if i == idx-1 { return iter.cur } iter.curPushNextStack(iter.cur) } else { return nil } } return cur } func (iter *vbtIterator) Next() (result bool) { if iter.dir > -1 { if iter.dir == 1 && iter.cur != nil { iter.tstack.Clear() iter.curPushNextStack(iter.cur) iter.up = iter.getNextUp(iter.cur) } iter.dir = -1 } if iter.tstack.Size() == 0 { if iter.up == nil { return false } iter.tstack.Push(iter.up) iter.up = iter.getNextUp(iter.up) } if v, ok := iter.tstack.Pop(); ok { iter.cur = v.(*tNode) iter.curPushNextStack(iter.cur) return true } return false } func (iter *vbtIterator) GetPrev(cur *tNode, idx int) *tNode { // iter := NewIterator(cur) iter.SetNode(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.(*tNode) if i == idx-1 { return iter.cur } iter.curPushPrevStack(iter.cur) } else { return nil } } return cur } func (iter *vbtIterator) Prev() (result bool) { if iter.dir < 1 { // 非 1(next 方向定义 -1 为 prev) if iter.dir == -1 && iter.cur != nil { // 如果上次为prev方向, 则清空辅助计算的栈 iter.tstack.Clear() iter.curPushPrevStack(iter.cur) // 把当前cur计算的逆向回朔 iter.up = iter.getPrevUp(iter.cur) // cur 寻找下个要计算up } iter.dir = 1 } // 如果栈空了, 把up的递归计算入栈, 重新计算 下次的up值 if iter.tstack.Size() == 0 { if iter.up == nil { return false } iter.tstack.Push(iter.up) iter.up = iter.getPrevUp(iter.up) } if v, ok := iter.tstack.Pop(); ok { iter.cur = v.(*tNode) iter.curPushPrevStack(iter.cur) return true } // 如果再次计算的栈为空, 则只能返回false return false } func getRelationship(cur *tNode) int { if cur.parent.children[1] == cur { return 1 } return 0 } func (iter *vbtIterator) getPrevUp(cur *tNode) *tNode { for cur.parent != nil { if getRelationship(cur) == 1 { // next 在 降序 小值. 如果child在右边, parent 比 child 小, parent才有效, 符合降序 return cur.parent } cur = cur.parent } return nil } func (iter *vbtIterator) curPushPrevStack(cur *tNode) { Prev := cur.children[0] // 当前的左然后向右找, 找到最大, 就是最接近cur 并且小于cur的值 if Prev != nil { iter.tstack.Push(Prev) for Prev.children[1] != nil { Prev = Prev.children[1] iter.tstack.Push(Prev) // 入栈 用于回溯 } } } func (iter *vbtIterator) getNextUp(cur *tNode) *tNode { for cur.parent != nil { if getRelationship(cur) == 0 { // Prev 在 降序 大值. 如果child在左边, parent 比 child 大, parent才有效 , 符合降序 return cur.parent } cur = cur.parent } return nil } func (iter *vbtIterator) curPushNextStack(cur *tNode) { next := cur.children[1] if next != nil { iter.tstack.Push(next) for next.children[0] != nil { next = next.children[0] iter.tstack.Push(next) } } }