bst_level_order.go

// https://leetcode.com/problems/binary-tree-level-order-traversal

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

type LevelNode struct {
    node *TreeNode
    level int
}

func levelOrder(root *TreeNode) [][]int {
    if root == nil {
        return [][]int{}
    }
    
    treeHeight := getTreeHeight(root)
    levels := make([][]int, treeHeight)
    queue := []*LevelNode{&LevelNode{node: root}}
    
    for len(queue) > 0 {
        curr := queue[0]
        queue = queue[1:]
        
        levels[curr.level] = append(levels[curr.level], curr.node.Val)
        
        if curr.node.Left != nil {
            queue = append(queue, &LevelNode{curr.node.Left, curr.level + 1})
        }
        
        if curr.node.Right != nil {
            queue = append(queue, &LevelNode{curr.node.Right, curr.level + 1})
        }
    }
    
    return levels
}

func getTreeHeight(root *TreeNode) int {
    if root == nil {
        return 0
    }
    
    return max(getTreeHeight(root.Left), getTreeHeight(root.Right)) + 1
}

func max(a, b int) int {
    if a < b {
        return b
    }
    
    return a
}