Skip to content

Instantly share code, notes, and snippets.

@cabloo
Created Jun 19, 2021
Embed
What would you like to do?
/**
* 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
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment