leetcode day 30 – I am not friends with slices: level-order tree

Visit the nodes of a binary tree in level order: all nodes at the root level, left-to-right, all nodes at level 1, left-to-right, etc.

I did try to solve this with slices and boy did I have trouble. I eventually stopped trying to use them and went to a list-oriented solution with a double closure.

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

type MyList struct {
    head *MyNode
    tail *MyNode
}

type MyNode struct {
    Val int
    Next *MyNode
}

func levelOrder(root *TreeNode) [][]int {
    // Ideally, I want to end up with an array
    // of arrays of nodes, added in left-to-right order.
    getStash, helper := createClosures()
    helper(root, 0)
    stash := getStash()

    out := [][]int{}
    for i := 0; i < len(stash); i++ {
        out = append(out, []int{})
        thisList := stash[i]
        for scan := thisList.head; scan != nil; scan = scan.Next {
            out[i] = append(out[i], scan.Val)
        }
    }
    return out
}

func createClosures() (func() []MyList, func(*TreeNode, int)) {
    stash := []MyList{}

    var helper func(*TreeNode, int)
    helper = func(root *TreeNode, level int) {
            if root == nil {
                // Nothing to do at this level
                return
            }
            // Current node gets stashed at the end of the list for this level.
            // (*output)[level] is the slice to append to.
            // Add new node to list at this level
            if len(stash) <= level {
                stash = append(stash, MyList{})
            }
            
            n := &MyNode{Val: root.Val}
            if stash[level].head == nil {
                stash[level].head = n
                stash[level].tail = n
            } else {
                stash[level].tail.Next = n
                stash[level].tail = n
            }
 
            // add the left and right subtrees at the next level down
            helper(root.Left, level + 1)
            helper(root.Right, level + 1)
        }
    return func() []MyList { return stash }, helper
}

This solves the problem, but it’s poor on memory (8th percentile) and runtime (26%). On the other hand, the list manipulation worked correctly the first time. Let’s try to speed it up. First let’s try to remove the lists.

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

func levelOrder(root *TreeNode) [][]int {
    // Ideally, I want to end up with an array
    // of arrays of nodes, added in left-to-right order.
    getStash, helper := createClosures()
    helper(root, 0)
    return getStash()
}

func createClosures() (func() [][]int, func(*TreeNode, int)) {
    // This will be our output structure. As we go down the
    // levels, we add more []int slices to hold the next level
    // of nodes. Because we work left to right, we always add
    // new nodes at a level to the correct slice, and they always
    // are added to the right end of the slice, giving us the
    // output data structure we want.
    stash := [][]int{}

    // We have to declare a named variable to be closed over for
    // a Go anonymous function to be able to call itself.
    var helper func(*TreeNode, int)

    // The real meat of the process.
    helper = func(root *TreeNode, level int) {
            if root == nil {
                // Nothing to do at this level
                return
            }
            // Current node gets stashed at the end of the slice
            //  for this level.
            // stash[level] is the slice to append to.
            if len(stash) <= level {
                // We have never accessed this level. Add a slice
                // so appends will work. You CANNOT just assign to
                // stash[level], because it doesn't exist yet and
                // you'll get an out of bounds error.
                stash = append(stash, []int{})
            }
            // Add this node's value to the end of the list at this
            // level.
            stash[level] = append(stash[level], root.Val)

            // add the left and right subtrees at the next level down.
            // Because we're tracking the level, we always append to the
            // right slice in the stash, which lives outside the recursion.
            helper(root.Left, level + 1)
            helper(root.Right, level + 1)
        }
    // The two functions the main function needs to call. Stash will be
    // exactly the needed data structure when helper finishes.
    return func() [][]int { return stash }, helper
}

That’s a significant improvement; 71st percentile runtime, 24th percentile memory. The problem I had the first time around was not properly managing the slices. To add an element to a slice after the current end, you must use append(), and I needed to do this for each level as I got there. The code to build the output structure from the lists was the learning experience I needed to get that part right in the rewrite.

Still definitely more comfortable manipulating data structures. Let’s try removing the closure and see how that goes; we’ll just let scoping make the stash visible to the helper.

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

func levelOrder(root *TreeNode) [][]int {
    // Ideally, I want to end up with an array
    // of arrays of nodes, added in left-to-right order.
    stash := [][]int{}

var helper func(*TreeNode, int)
    helper = func(root *TreeNode, level int) {
            if root == nil {
                // Nothing to do at this level
                return
            }
            // Current node gets stashed at the end of the list for this level.
            // stash[level] is the slice to append to.
            // Add new node to list at this level
            if len(stash) <= level {
                stash = append(stash, []int{})
            }
            stash[level] = append(stash[level], root.Val)

            // add the left and right subtrees at the next level down
            helper(root.Left, level + 1)
            helper(root.Right, level + 1)
        }

    helper(root, 0)
    return stash
}

Interestingly, this is actually worse according to leetcode; 68th percentile now (3ms instead of 1ms). Honestly I think we’re in the noise at this point. No change in the memory usage.

I think if I’d gotten this in an interview I’d call the second or third solution good, and the first marginal. The rewrite time from the list version to the working slice version was only a couple minutes, and I never felt like I didn’t grasp the necessities of the problem.

Reply