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
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
|
// 2458. Height of Binary Tree After Subtree Removal Queries
#include "leetcode.h"
/*
* 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
*/
#define MAX_NODES 100001
struct node_info {
int height;
int level;
};
struct level_info {
int max_height1;
int max_height2;
};
struct node_info *node_map[MAX_NODES];
struct level_info *level_map[MAX_NODES];
int dfs(struct TreeNode *root, int level) {
if (!root)
return 0;
int left = dfs(root->left, level + 1);
int right = dfs(root->right, level + 1);
int height = left > right ? left : right;
if (!level_map[level]) {
level_map[level] = (struct level_info *)malloc(sizeof(struct level_info));
level_map[level]->max_height1 = -1;
level_map[level]->max_height2 = -1;
}
if (height >= level_map[level]->max_height1) {
level_map[level]->max_height2 = level_map[level]->max_height1;
level_map[level]->max_height1 = height;
} else if (height > level_map[level]->max_height2)
level_map[level]->max_height2 = height;
node_map[root->val] = (struct node_info *)malloc(sizeof(struct node_info));
node_map[root->val]->height = height;
node_map[root->val]->level = level;
return height + 1;
}
int *treeQueries(struct TreeNode *root, int *queries, int queriesSize,
int *returnSize) {
*returnSize = queriesSize;
int max_depth = dfs(root, 0) - 1;
int *ans = (int *)malloc(queriesSize * sizeof(int));
for (int i = 0; i < queriesSize; i++) {
int curr = queries[i];
struct node_info *n = node_map[curr];
int height = n->height;
int level = n->level;
struct level_info *l = level_map[level];
if (l->max_height1 == height)
ans[i] = l->max_height2 != -1
? max_depth - abs(l->max_height1 - l->max_height2)
: max_depth - 1;
else
ans[i] = max_depth;
}
for (int i = 0; i < MAX_NODES; i++) {
if (node_map[i]) {
free(node_map[i]);
node_map[i] = NULL;
}
if (level_map[i]) {
free(level_map[i]);
level_map[i] = NULL;
}
}
return ans;
}
int main() {
struct TreeNode *root = treenode_create(5);
root->left = treenode_create(8);
root->right = treenode_create(9);
root->left->left = treenode_create(2);
root->left->right = treenode_create(1);
root->right->left = treenode_create(3);
root->right->right = treenode_create(7);
root->left->left->left = treenode_create(4);
root->left->left->right = treenode_create(6);
int queries[] = {3, 2, 4, 8}, returnSize;
int *tq = treeQueries(root, queries, ARRAY_SIZE(queries), &returnSize);
for (int i = 0; i < returnSize; i++)
printf("%d ", tq[i]); // expect: 3 2 3 2
printf("\n");
free(tq);
treenode_free(root);
}
|