leetcode-go/leetcode/medium/102.go

package main

// MyNode ...
type MyNode struct {
	Node  *TreeNode
	Level int
}

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func levelOrder(root *TreeNode) [][]int {
	result := [][]int{}
	stack := []MyNode{}
	curr := root
	top := -1
	level := 0
	for curr != nil || top > -1 { // Inorder traversal
		for curr != nil {
			stack = append(stack, MyNode{curr, level})
			top++
			curr = curr.Left
			level++
		}
		if top > -1 {
			curr = stack[top].Node
			level = stack[top].Level
			stack = stack[:top]
			top--
			for len(result) <= level {
				result = append(result, []int{})
			}
			result[level] = append(result[level], curr.Val)
			curr = curr.Right
			level++
		}
	}
	return result
}

// func main() {
// 	//        1
// 	//       /
// 	//      2
// 	//     / \
// 	//    3   4
// 	n4 := TreeNode{4, nil, nil}
// 	n3 := TreeNode{3, nil, nil}
// 	n2 := TreeNode{2, &n3, &n4}
// 	n1 := TreeNode{1, &n2, nil}
// 	fmt.Println(levelOrder(&n1))
// }