Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
/**
* 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
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment