bst_zigzag_order.go

// https://leetcode.com/problems/binary-tree-zigzag-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 zigzagLevelOrder(root *TreeNode) [][]int {
    if root == nil {
        return [][]int{}
    }
    
    treeHeight := getTreeHeight(root)
    result := make([][]int, treeHeight)
    queue := []*LevelNode{&LevelNode{
        node: root,
        level: 0,
    }}
    
    for len(queue) > 0 {
        curr := queue[0]
        queue = queue[1:]
        
        result[curr.level] = append(result[curr.level], curr.node.Val)
        
        if curr.node.Left != nil {
            queue = append(queue, &LevelNode{
                node: curr.node.Left,
                level: curr.level + 1,
            })
        }
        
        if curr.node.Right != nil {
            queue = append(queue, &LevelNode{
                node: curr.node.Right,
                level: curr.level + 1,
            })
        }
    }
    
    for level, items := range result {
        if level & 1 == 1 {
            reverse(items)
        }
    }
    
    return result
}

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

func reverse(arr []int) {
    for i, j := 0, len(arr) - 1; i < j; i, j = i + 1, j - 1 {
        arr[i], arr[j] = arr[j], arr[i]
    }
}

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