LeetCode in Kotlin

814. Binary Tree Pruning

Medium

Given the root of a binary tree, return the same tree where every subtree (of the given tree) not containing a 1 has been removed.

A subtree of a node node is node plus every node that is a descendant of node.

Example 1:

Input: root = [1,null,0,0,1]

Output: [1,null,0,null,1]

Explanation:

Only the red nodes satisfy the property “every subtree not containing a 1”.

The diagram on the right represents the answer.

Example 2:

Input: root = [1,0,1,0,0,0,1]

Output: [1,null,1,null,1]

Example 3:

Input: root = [1,1,0,1,1,0,1,0]

Output: [1,1,0,1,1,null,1]

Constraints:

Solution

import com_github_leetcode.TreeNode

/*
 * Example:
 * var ti = TreeNode(5)
 * var v = ti.`val`
 * Definition for a binary tree node.
 * class TreeNode(var `val`: Int) {
 *     var left: TreeNode? = null
 *     var right: TreeNode? = null
 * }
 */
class Solution {
    fun pruneTree(root: TreeNode?): TreeNode? {
        if (root == null) {
            return root
        }
        root.left = pruneTree(root.left)
        root.right = pruneTree(root.right)
        return if (root.left == null && root.right == null && root.`val` == 0) {
            null
        } else {
            root
        }
    }
}