btreePrune.go

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

type state struct {
    parent *TreeNode
    node *TreeNode
    haveTraversed bool
}

func popStack(stack []*state) (*state, []*state) {
    end := len(stack)-1
    curr := stack[end]
    stack[end] = nil
    stack = stack[:end]
    return curr, stack
}

func pruneTree(root *TreeNode) *TreeNode {
    stack := []*state{{nil, root, false}}
    for len(stack) > 0 {
        var curr *state
        curr, stack = popStack(stack)
        if curr == nil || curr.node == nil {
            continue
        }
        node := curr.node
        
        if !curr.haveTraversed && (node.Left != nil || node.Right != nil) {
            stack = append(stack, []*state{
                {curr.parent, node, true},
                {node, node.Left, false},
                {node, node.Right, false},
            }...)
            continue
        }
        
        if node.Val == 0 && node.Left == nil && node.Right == nil {
            if curr.parent == nil {
                return nil
            }
            if curr.parent.Left == node {
                curr.parent.Left = nil
                continue
            }
            curr.parent.Right = nil
        }
    }

    return root
}