constructBTree.go

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

type stateEntry struct {
    inOrderOffset int
    preOrder, inOrder []int
    top *TreeNode
}

type stateStack struct {
    items []*stateEntry
}

func (s *stateStack) Pop() *stateEntry {
    if len(s.items) == 0 {
        return nil
    }
    
    end := len(s.items)-1
    state := s.items[end]
    s.items[end] = nil
    s.items = s.items[:end]
    
    return state
}

func (s *stateStack) Push(state *stateEntry) {
    s.items = append(s.items, state)
}

func buildTree(preorder []int, inorder []int) *TreeNode {
    if len(preorder) != len(inorder) || len(preorder) == 0 {
        return nil
    }
    
    whatIndexInInOrder := map[int]int{}
    for i, v := range inorder {
        whatIndexInInOrder[v] = i
    }

    top := &TreeNode{Val: preorder[0]}
    states := &stateStack{}
    states.Push(&stateEntry{
        top: top,
        preOrder: preorder,
        inOrder: inorder,
    })

    for state := states.Pop(); state != nil; state = states.Pop() {
        leftLength := whatIndexInInOrder[state.top.Val] - state.inOrderOffset
        rightLength := len(state.inOrder) - leftLength - 1
        
        if leftLength > 0 {
            state.top.Left = &TreeNode{Val: state.preOrder[1]}
        }
        
        if leftLength > 1 {
            states.Push(&stateEntry{
                inOrderOffset: state.inOrderOffset,
                inOrder: state.inOrder[:leftLength],
                preOrder: state.preOrder[1:leftLength+1],
                top: state.top.Left,
            })
        }

        if rightLength > 0 {
            state.top.Right = &TreeNode{Val: state.preOrder[leftLength+1]}
        }

        if rightLength > 1 {
            states.Push(&stateEntry{
                inOrderOffset: state.inOrderOffset + leftLength + 1,
                inOrder: state.inOrder[leftLength+1:],
                preOrder: state.preOrder[leftLength+1:],
                top: state.top.Right,
            })
        }
    }

    return top
}