【二叉树 C++DFS】2458. 移除子树后的二叉树高度

本文涉及知识点

二叉树 C++DFS

LeetCode 2458. 移除子树后的二叉树高度

给你一棵 二叉树 的根节点 root ，树中有 n 个节点。每个节点都可以被分配一个从 1 到 n 且互不相同的值。另给你一个长度为 m 的数组 queries 。
你必须在树上执行 m 个 独立 的查询，其中第 i 个查询你需要执行以下操作：
从树中 移除 以 queries[i] 的值作为根节点的子树。题目所用测试用例保证 queries[i] 不 等于根节点的值。
返回一个长度为 m 的数组 answer ，其中 answer[i] 是执行第 i 个查询后树的高度。
注意：
查询之间是独立的，所以在每个查询执行后，树会回到其 初始 状态。
树的高度是从根到树中某个节点的 最长简单路径中的边数 。
示例 1：
输入：root = [1,3,4,2,null,6,5,null,null,null,null,null,7], queries = [4]
输出：[2]
解释：上图展示了从树中移除以 4 为根节点的子树。
树的高度是 2（路径为 1 -> 3 -> 2）。
示例 2：
输入：root = [5,8,9,2,1,3,7,4,6], queries = [3,2,4,8]
输出：[3,2,3,2]
解释：执行下述查询：

• 移除以 3 为根节点的子树。树的高度变为 3（路径为 5 -> 8 -> 2 -> 4）。
• 移除以 2 为根节点的子树。树的高度变为 2（路径为 5 -> 8 -> 1）。
• 移除以 4 为根节点的子树。树的高度变为 3（路径为 5 -> 8 -> 2 -> 6）。
• 移除以 8 为根节点的子树。树的高度变为 2（路径为 5 -> 9 -> 3）。
  提示：
  树中节点的数目是 n
  2 <= n <= 10⁵
  1 <= Node.val <= n
  树中的所有值 互不相同
  m == queries.length
  1 <= m <= min(n, 10⁴)
  1 <= queries[i] <= n
  queries[i] != root.val

深度优先搜索(错误）

DFS的参数：TreeNode* root ,int height 。height是当前节点到根节点的最短距离。
DFS的返回值：当前子树的高度。
m_vRet[i] 记录删除值为i的子树后，高度减少多少。
DFS：过程。
记录各子树的高度。
如果没有子树，返回。
令child1是最高的子树，child2是次高的子树。
m_vRet[child1] = child1的高度-child2的高度。
如果child2不存在，可以假定它的高度为-1。
对个查询的值，整棵树的高度-m_vRet[que]
错误原因： 只有删除同层次(leve)节点高度最高的子树时，高度才会发生变化。不是兄弟节点的最高，次高；而是整个层次的最高次高。

深度优先搜索(更正）

DFS的参数：TreeNode* cur,int leve
DFS返回值：m_leve2[cur->value]
DFS：过程
m_leve[cur]记录，cur的层次，根节点是0。
m_leve2[cur]记录，cur子树的层次数量。 叶子节点为1。
分如下几步：
int iHeight = DFS(root,0);
leveToLeve2 记录各层次的leve2
对leveToLeve2[i]按降序排序
计算各查询的sub:
auto& v = leveToLeve2[m_leve[que]];
如果不是当前层次的最大leve2，则sub为0。
如果当前层次只有一个节点，sub = v[0]
否则 sub = v[1]
当前查询的值就是：iHeight-1 - sub。
时间复杂度：O(nlogn) 理论上瓶颈在排序，实际上在DFS。力扣上的DFS非常花时间。

代码

核心代码

#define MaxValue (100'000)
class Solution {
public:
	vector<int> treeQueries(TreeNode* root, vector<int>& queries) {
		int iHeight = DFS(root,0);
		vector<int> leveToLeve2[MaxValue + 1];
		for (int i = 1; i <= MaxValue; i++) {
			leveToLeve2[m_leve[i]].emplace_back(m_leve2[i]);
		}
		for (int i = 0; i < MaxValue; i++) {
			sort(leveToLeve2[i].begin(), leveToLeve2[i].end(),greater<>());
		}
		vector<int> ret;
		for (const auto& que : queries) {
			int sub = 0;
			auto& v = leveToLeve2[m_leve[que]];
			if (m_leve2[que] == v[0]) {
				sub = (1 == v.size()) ? v[0] : (v[0] - v[1]);
			}
			ret.emplace_back(iHeight-1 - sub);
		}
		return ret;
	}
	int DFS(TreeNode* cur,int leve) {
		if (nullptr == cur) { return 0; }
		m_leve[cur->val] = leve;
		const int i1 = DFS(cur->left,leve+1);
		const int i2 = DFS(cur->right,leve+1);		
		return m_leve2[cur->val] = max(i1, i2) + 1;
	}
	int m_leve[MaxValue + 1] = { 0 };
	int m_leve2[MaxValue + 1] = { 0 };
	
};
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

单元测试

template<class T1, class T2>
void AssertEx(const T1& t1, const T2& t2)
{
	Assert::AreEqual(t1, t2);
}

template<class T>
void AssertEx(const vector<T>& v1, const vector<T>& v2)
{
	Assert::AreEqual(v1.size(), v2.size());
	for (int i = 0; i < v1.size(); i++)
	{
		Assert::AreEqual(v1[i], v2[i]);
	}
}

template<class T>
void AssertV2(vector<vector<T>> vv1, vector<vector<T>> vv2)
{
	sort(vv1.begin(), vv1.end());
	sort(vv2.begin(), vv2.end());
	Assert::AreEqual(vv1.size(), vv2.size());
	for (int i = 0; i < vv1.size(); i++)
	{
		AssertEx(vv1[i], vv2[i]);
	}
}

namespace UnitTest
{
	vector<int> nums, queries;
	const int null = 100'000;
	TEST_CLASS(UnitTest)
	{
	public:
		TEST_METHOD(TestMethod00)
		{
			nums = { 1, 3, 4, 2, null, 6, 5, null, null, null, null, null, 7 }, queries = { 4 };
			auto root = NTree::Init(nums,null);
			auto res = Solution().treeQueries(root, queries);
			AssertEx(vector<int>{2}, res);
		}
		TEST_METHOD(TestMethod01)
		{
			nums = { 5,8,9,2,1,3,7,4,6 }, queries = { 3,2,4,8 };
			auto root = NTree::Init(nums, null);
			auto res = Solution().treeQueries(root, queries);
			AssertEx(vector<int>{3, 2, 3, 2}, res);
		}
		
	};
}
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

求第k小函数优化

nth_element(a,a+k,a+n);
功能：函数只是把下标为k的元素放在了正确位置，对其它元素并没有排序当然k左边元素都小于等于它，右边元素都大于等于它，所以可以利用这个函数快速定位某个元素。
理论上可以把时间复杂度降到：O(n) 实际上稍稍提高速度。

class Solution {
public:
	vector<int> treeQueries(TreeNode* root, vector<int>& queries) {
		int iHeight = DFS(root,0);
		vector<int> leveToLeve2[MaxValue + 1];
		for (int i = 1; i <= MaxValue; i++) {
			leveToLeve2[m_leve[i]].emplace_back(m_leve2[i]);
		}
		for (int i = 0; i < MaxValue; i++) {
			if(leveToLeve2[i].size() >= 2 )
			nth_element(leveToLeve2[i].begin(), leveToLeve2[i].end()-2,leveToLeve2[i].end());
		}
		vector<int> ret;
		for (const auto& que : queries) {
			int sub = 0;
			auto& v = leveToLeve2[m_leve[que]];
			if (m_leve2[que] == v.back()) {
				sub = (1 == v.size()) ? v.back() : (v.back() - v[v.size()-2]);
			}
			ret.emplace_back(iHeight-1 - sub);
		}
		return ret;
	}
	int DFS(TreeNode* cur,int leve) {
		if (nullptr == cur) { return 0; }
		m_leve[cur->val] = leve;
		const int i1 = DFS(cur->left,leve+1);
		const int i2 = DFS(cur->right,leve+1);		
		return m_leve2[cur->val] = max(i1, i2) + 1;
	}
	int m_leve[MaxValue + 1] = { 0 };
	int m_leve2[MaxValue + 1] = { 0 };
	
};
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

扩展阅读

视频课程

先学简单的课程，请移步CSDN学院，听白银讲师（也就是鄙人）的讲解。
https://edu.csdn.net/course/detail/38771

如何你想快速形成战斗了，为老板分忧，请学习C#入职培训、C++入职培训等课程
https://edu.csdn.net/lecturer/6176

相关推荐

测试环境

操作系统：win7 开发环境： VS2019 C++17
或者 操作系统：win10 开发环境： VS2022 C++17
如无特殊说明，本算法用**C++**实现。