cabloo · June 20, 2021 13:51 · Jun 20, 2021
diff --git a/constructBTree.go b/constructBTree.go
@@ -0,0 +1,87 @@
+/**
+ * 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
+}