65 lines
1.2 KiB
Go
65 lines
1.2 KiB
Go
package q173
|
|
|
|
type TreeNode struct {
|
|
Val int
|
|
Left *TreeNode
|
|
Right *TreeNode
|
|
}
|
|
|
|
type stackElem struct {
|
|
state uint8 // 0, 1, 2, 3
|
|
node *TreeNode
|
|
}
|
|
|
|
type BSTIterator struct {
|
|
stack []stackElem
|
|
next *TreeNode
|
|
}
|
|
|
|
func Constructor(root *TreeNode) BSTIterator {
|
|
return BSTIterator{stack: []stackElem{{node: root}}}
|
|
}
|
|
|
|
func (it *BSTIterator) Next() int {
|
|
it.findNext()
|
|
ret := it.next.Val
|
|
it.next = nil
|
|
return ret
|
|
}
|
|
|
|
func (it *BSTIterator) HasNext() bool {
|
|
it.findNext()
|
|
return it.next != nil
|
|
}
|
|
|
|
func (it *BSTIterator) findNext() {
|
|
for it.next == nil && len(it.stack) > 0 {
|
|
c := &it.stack[len(it.stack)-1]
|
|
c.state++
|
|
|
|
switch c.state - 1 {
|
|
case 0: // init
|
|
if c.node.Left != nil {
|
|
it.stack = append(it.stack, stackElem{node: c.node.Left})
|
|
}
|
|
|
|
case 1: // left processed
|
|
it.next = c.node
|
|
|
|
case 2: // middle processed
|
|
if c.node.Right != nil {
|
|
it.stack = append(it.stack, stackElem{node: c.node.Right})
|
|
}
|
|
|
|
case 3: // right processed
|
|
it.stack = it.stack[:len(it.stack)-1]
|
|
}
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Your BSTIterator object will be instantiated and called as such:
|
|
* obj := Constructor(root);
|
|
* param_1 := obj.Next();
|
|
* param_2 := obj.HasNext();
|
|
*/
|