diff options
author | Jack Sangdahl <[email protected]> | 2024-10-26 00:09:07 -0600 |
---|---|---|
committer | Jack Sangdahl <[email protected]> | 2024-10-26 00:09:07 -0600 |
commit | a43f21d0aff66470e17b47febcd01e67bbdce5f6 (patch) | |
tree | 56f0a0d0f02e8aa79e6a1426f8ba128552404a09 /height_binary_tree_removal.py | |
parent | 54e0305c7a0927ea74163843e547dcb93d84344c (diff) | |
download | leetcode-a43f21d0aff66470e17b47febcd01e67bbdce5f6.tar.gz leetcode-a43f21d0aff66470e17b47febcd01e67bbdce5f6.zip |
2458. height of binarytree after subtree removal queries
Diffstat (limited to 'height_binary_tree_removal.py')
-rw-r--r-- | height_binary_tree_removal.py | 66 |
1 files changed, 66 insertions, 0 deletions
diff --git a/height_binary_tree_removal.py b/height_binary_tree_removal.py new file mode 100644 index 0000000..4ad028f --- /dev/null +++ b/height_binary_tree_removal.py @@ -0,0 +1,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])) |