Leetcode Problem 1110 - Delete Nodes And Return Forest.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Problem Statement. <br/>\n",
    "\n",
    "Given the root of a binary tree, each node in the tree has a distinct value. <br/>\n",
    "After deleting all nodes with a value in to_delete, we are left with a forest (a disjoint union of trees). <br/>\n",
    "Return the roots of the trees in the remaining forest.  You may return the result in any order. <br/>\n",
    "\n",
    "Example 1: <br/>\n",
    "Input: root = [1,2,3,4,5,6,7], to_delete = [3,5] <br/>\n",
    "Output: [[1,2,null,4],[6],[7]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BFS - O(N) runtime, O(N + D) space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "\n",
    "class Solution:\n",
    "    def delNodes(self, root: TreeNode, to_delete: List[int]) -> List[TreeNode]:\n",
    "        \n",
    "        res = set()\n",
    "        queue = deque([root])\n",
    "        to_delete_set = set(to_delete)\n",
    "        if root.val in to_delete_set:\n",
    "            if root.left: res.add(root.left)\n",
    "            if root.right: res.add(root.right)\n",
    "        else:\n",
    "            res.add(root)\n",
    "        \n",
    "        while queue:\n",
    "            node = queue.popleft()\n",
    "            if node.left: \n",
    "                queue.append(node.left)\n",
    "                if node.left.val in to_delete_set:\n",
    "                    if node.left.left: res.add(node.left.left)\n",
    "                    if node.left.right: res.add(node.left.right)\n",
    "                    node.left = None\n",
    "            if node.right: \n",
    "                queue.append(node.right)\n",
    "                if node.right.val in to_delete_set:\n",
    "                    if node.right.left: res.add(node.right.left)\n",
    "                    if node.right.right: res.add(node.right.right)\n",
    "                    node.right = None\n",
    "                        \n",
    "            if node.val in to_delete_set and node in res:\n",
    "                res.remove(node)\n",
    "                \n",
    "        return list(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DFS Stack - O(N) runtime, O(N + D) space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "#Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "\n",
    "class Solution:\n",
    "    def delNodes(self, root: TreeNode, to_delete: List[int]) -> List[TreeNode]:\n",
    "        \n",
    "        res = set()\n",
    "        stack = deque([root])\n",
    "        to_delete_set = set(to_delete)\n",
    "        if root.val in to_delete_set:\n",
    "            if root.left: res.add(root.left)\n",
    "            if root.right: res.add(root.right)\n",
    "        else:\n",
    "            res.add(root)\n",
    "        \n",
    "        while stack:\n",
    "            node = stack.pop()\n",
    "            if node.left: \n",
    "                stack.append(node.left)\n",
    "                if node.left.val in to_delete_set:\n",
    "                    if node.left.left: res.add(node.left.left)\n",
    "                    if node.left.right: res.add(node.left.right)\n",
    "                    node.left = None\n",
    "            if node.right: \n",
    "                stack.append(node.right)\n",
    "                if node.right.val in to_delete_set:\n",
    "                    if node.right.left: res.add(node.right.left)\n",
    "                    if node.right.right: res.add(node.right.right)\n",
    "                    node.right = None\n",
    "                        \n",
    "            if node.val in to_delete_set and node in res:\n",
    "                res.remove(node)\n",
    "                \n",
    "        return list(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DFS Recursive - O(N) runtime, O(N + D) space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List, Set\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "\n",
    "class Solution:\n",
    "    def delNodes(self, root: TreeNode, to_delete: List[int]) -> List[TreeNode]:\n",
    "        \n",
    "        res = set()\n",
    "        to_delete_set = set(to_delete)\n",
    "        if root.val in to_delete_set:\n",
    "            if root.left: res.add(root.left)\n",
    "            if root.right: res.add(root.right)\n",
    "        else:\n",
    "            res.add(root)\n",
    "            \n",
    "        self.delNodesRecursive(root, res, to_delete_set)\n",
    "        \n",
    "        return list(res)\n",
    "        \n",
    "    def delNodesRecursive(self, node: TreeNode, res: Set[TreeNode], to_delete_set: Set[int]) -> List[TreeNode]:\n",
    "        if node.left: \n",
    "            self.delNodesRecursive(node.left, res, to_delete_set)\n",
    "            if node.left.val in to_delete_set:\n",
    "                if node.left.left: res.add(node.left.left)\n",
    "                if node.left.right: res.add(node.left.right)\n",
    "                node.left = None         \n",
    "        \n",
    "        if node.right: \n",
    "            self.delNodesRecursive(node.right, res, to_delete_set)   \n",
    "            if node.right.val in to_delete_set:\n",
    "                if node.right.left: res.add(node.right.left)\n",
    "                if node.right.right: res.add(node.right.right)\n",
    "                node.right = None\n",
    "\n",
    "        if node.val in to_delete_set and node in res:\n",
    "            res.remove(node)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cleaner DFS Recursive - O(N) runtime, O(N + H) space where H is the height of the tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "\n",
    "class Solution:\n",
    "    def delNodes(self, root: TreeNode, to_delete: List[int]) -> List[TreeNode]:\n",
    "        to_delete_set = set(to_delete)\n",
    "        res = []\n",
    "\n",
    "        def helper(root, is_root):\n",
    "            if not root: return None\n",
    "            root_deleted = root.val in to_delete_set\n",
    "            if is_root and not root_deleted:\n",
    "                res.append(root)\n",
    "            root.left = helper(root.left, root_deleted)\n",
    "            root.right = helper(root.right, root_deleted)\n",
    "            return None if root_deleted else root\n",
    "        helper(root, True)\n",
    "        return res"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}