leetcode-go/medium/654.maximum-binary-tree.go

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

type maxHeap []int

func (h maxHeap) Len() int           { return len(h) }
func (h maxHeap) Less(i, j int) bool { return h[j] < h[i] }
func (h maxHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *maxHeap) Push(x interface{}) {
	*h = append(*h, x.(int))
}

func (h *maxHeap) Pop() interface{} {
	l := h.Len()
	x := (*h)[l-1]
	*h = (*h)[:l-1]
	return x
}

func constructMaximumBinaryTree(nums []int) *TreeNode {
	if len(nums) == 0 {
		return (*TreeNode)(nil)
	}
	m := make(map[int]int)
	var h maxHeap
	for i, v := range nums {
		m[v] = i
		heap.Push(&h, v)
	}
	root := TreeNode{heap.Pop(&h).(int), nil, nil}
	for h.Len() != 0 {
		curr := &root
		v := heap.Pop(&h).(int)
		i := m[v]
		for {
			j := m[curr.Val]
			if i < j {
				if curr.Left == nil {
					curr.Left = &TreeNode{v, nil, nil}
					break
				}
				curr = curr.Left
			} else {
				if curr.Right == nil {
					curr.Right = &TreeNode{v, nil, nil}
					break
				}
				curr = curr.Right
			}
		}
	}
	return &root
}