/** * 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 }