binary-search-tree-iterator.go

// https://leetcode.com/problems/binary-search-tree-iterator

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
type BSTIterator struct {
    stack []*TreeNode
}

func Constructor(root *TreeNode) BSTIterator {
    iterator := BSTIterator{
        stack: []*TreeNode{},
    }
    
    iterator.walkToMin(root)
    
    return iterator
}

// Next returns the next smallest number
func (i *BSTIterator) Next() int {
    value := -1
    stackSize := len(i.stack)
    
    if stackSize > 0 {
        current := i.stack[stackSize - 1]
        i.stack = i.stack[:stackSize - 1]
        value = current.Val
        
        i.walkToMin(current.Right)
    }
    
    return value
}

// HasNext tells whether we have a next smallest number or not
func (i *BSTIterator) HasNext() bool {
    return len(i.stack) > 0
}

// walkToMin walks to the minimum element of the given subtree
func (i *BSTIterator) walkToMin(root *TreeNode) {
    for root != nil {
        i.stack = append(i.stack, root)
        root = root.Left
    }
}