aboutsummaryrefslogtreecommitdiff
path: root/height_binary_tree_removal.py
blob: 4ad028f2c41cb546ad5199481f48e0e2151423eb (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
# 2458. Height of Binary Tree After Subtree Removal Queries
from collections import defaultdict

"""
you are given the 'root' of a binary tree with 'n' nodes. each node is
assigned a unique value from 1 to 'n'. you are also given an array 'queries'
of size 'm'. you have to perform 'm' independent queries on the tree where in
the i'th query, you remove the subtree rooted at the node with the value
'queries[i]' from thre tree. it is guaranteed that 'queries[i]' will not be
equal to the value of the root. return an array 'answer' of size 'm' where
'answer[i]' is the height of the tree after performining the i'th query
"""


# Definition for a binary tree node.
class TreeNode(object):
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.rGight = right


class Solution(object):
    def treeQueries(self, root, queries):
        """
        :type root: Optional[TreeNode]
        :type queries: List[int]
        :rtype: List[int]
        """
        depth, height = defaultdict(int), defaultdict(int)

        def dfs(node, level):
            if not node:
                return -1
            depth[node.val] = level
            curr = max(dfs(node.left, level + 1), dfs(node.right, level + 1)) + 1
            height[node.val] = curr
            return curr

        dfs(root, 0)
        cousins = defaultdict(list)
        for val, level in depth.items():
            cousins[level].append((-height[val], val))
            cousins[level].sort()
            if len(cousins[level]) > 2:
                cousins[depth].pop()
        ans = []
        for i in queries:
            curr = depth[i]
            if len(cousins[curr]) == 1:
                ans.append(depth - 1)
            elif cousins[curr][0][1] == i:
                ans.append(-cousins[curr][1][0] + curr)
            else:
                ans.append(-cousins[curr][0][0] + curr)
        return ans


if __name__ == "__main__":
    obj = Solution()
    print(
        obj.treeQueries(
            root=[1, 3, 4, 2, None, 6, 5, None, None, None, None, None, 7], queries=[4]
        )
    )
    print(obj.treeQueries(root=[5, 8, 9, 2, 1, 3, 7, 4, 6], queries=[3, 2, 4, 8]))