【树】【图论】【树路径】【深度优先搜索】2867. 统计树中的合法路径数目

作者推荐

【广度优先搜索】【网格】【割点】【 推荐】1263. 推箱子

本文涉及知识点

树 图论 深度优先搜索

LeetCode:2867. 统计树中的合法路径数目

给你一棵 n 个节点的无向树，节点编号为 1 到 n 。给你一个整数 n 和一个长度为 n - 1 的二维整数数组 edges ，其中 edges[i] = [ui, vi] 表示节点 ui 和 vi 在树中有一条边。
请你返回树中的 合法路径数目 。
如果在节点 a 到节点 b 之间 恰好有一个 节点的编号是质数，那么我们称路径 (a, b) 是 合法的 。
注意：
路径 (a, b) 指的是一条从节点 a 开始到节点 b 结束的一个节点序列，序列中的节点 互不相同 ，且相邻节点之间在树上有一条边。
路径 (a, b) 和路径 (b, a) 视为 同一条 路径，且只计入答案 一次 。
示例 1：
输入：n = 5, edges = [[1,2],[1,3],[2,4],[2,5]]
输出：4
解释：恰好有一个质数编号的节点路径有：

• (1, 2) 因为路径 1 到 2 只包含一个质数 2 。
• (1, 3) 因为路径 1 到 3 只包含一个质数 3 。
• (1, 4) 因为路径 1 到 4 只包含一个质数 2 。
• (2, 4) 因为路径 2 到 4 只包含一个质数 2 。
  只有 4 条合法路径。
  示例 2：
  输入：n = 6, edges = [[1,2],[1,3],[2,4],[3,5],[3,6]]
  输出：6
  解释：恰好有一个质数编号的节点路径有：
• (1, 2) 因为路径 1 到 2 只包含一个质数 2 。
• (1, 3) 因为路径 1 到 3 只包含一个质数 3 。
• (1, 4) 因为路径 1 到 4 只包含一个质数 2 。
• (1, 6) 因为路径 1 到 6 只包含一个质数 3 。
• (2, 4) 因为路径 2 到 4 只包含一个质数 2 。
• (3, 6) 因为路径 3 到 6 只包含一个质数 3 。
  只有 6 条合法路径。
  提示：
  1 <= n <= 10⁵
  edges.length == n - 1
  edges[i].length == 2
  1 <= ui, vi <= n
  输入保证 edges 形成一棵合法的树。

深度优先搜索

以任意节点为根，任意两个节点的路径一定是： 起点 $\to$ 公共祖先 $\to$ 终点，特例：起点或终点就是公共祖先。我们枚举公共祖先。DFS返回两个值：子树根节点为起点，子树的任意节点为终点的路径数。第一个返回值：经过一个质数。第二个返回值：经过0个质数。

分情况讨论

一，子树根节点是非质数。
a,特例。各子树经过1质数的路径。
b,各子树0质数的路径和兄长1质数路径结合。
c,各个子数1质数路径和兄长0质数路径结合。
二，子树根节点是质数。
a,特例。各子树经过0质数的路径和。
b,各子树0质数的路径和兄长0质数路径结合。
总结：
左支：a，根节点。 b，根节点+兄长节点的某一路径。
右支：当前支。
总共两种情况：
一：左支1质数，右支0 质数。
二：左支0质数，右支1 质数。

代码

核心代码

class CNeiBo2
{
public:
	CNeiBo2(int n, bool bDirect, int iBase = 0) :m_iN(n), m_bDirect(bDirect), m_iBase(iBase)
	{
		m_vNeiB.resize(n);
	}
	CNeiBo2(int n, vector<vector<int>>& edges, bool bDirect, int iBase = 0) :m_iN(n), m_bDirect(bDirect), m_iBase(iBase)
	{
		m_vNeiB.resize(n);
		for (const auto& v : edges)
		{
			m_vNeiB[v[0] - iBase].emplace_back(v[1] - iBase);
			if (!bDirect)
			{
				m_vNeiB[v[1] - iBase].emplace_back(v[0] - iBase);
			}
		}
	}
	inline void Add(int iNode1, int iNode2)
	{
		iNode1 -= m_iBase;
		iNode2 -= m_iBase;
		m_vNeiB[iNode1].emplace_back(iNode2);
		if (!m_bDirect)
		{
			m_vNeiB[iNode2].emplace_back(iNode1);
		}
	}
	const int m_iN;
	const bool m_bDirect;
	const int m_iBase;
	vector<vector<int>> m_vNeiB;
};

vector<int> CreatePrime(int iMax)
{
	vector<int> vPrime = { 2 };
	for (int i = 3; i <= iMax; i++)
	{
		bool b = true;
		for (const auto& n : vPrime)
		{
			if (0 == i % n)
			{
				b = false;
				break;
			}
		}
		if (b)
		{
			vPrime.emplace_back(i);
		}
	}
	return vPrime;
}

class Solution {
public:
	long long countPaths(int n, vector<vector<int>>& edges) {
		auto v = CreatePrime(n);
		n_setPrime.insert(v.begin(), v.end());
		CNeiBo2 neiBo2(n, edges, false, 1);
		for (auto& v : neiBo2.m_vNeiB)
		{
			sort(v.begin(), v.end());
		}
		DFS(neiBo2.m_vNeiB, 0, -1);
		return m_llRet;
	}
	pair<long long, long long> DFS(vector<vector<int>>& neiBo, int cur, int par)
	{
		long long i0=0, i1=0;
		if (n_setPrime.count(cur+1))
		{
			i1 = 1;			
		}
		else
		{
			i0 = 1;
		}		
		for (const auto& next : neiBo[cur])
		{
			if (next == par)
			{
				continue;
			}
			
			const auto [i01, i11] = DFS(neiBo, next, cur);
			m_llRet += i11 * i0;
			m_llRet += i01 * i1;
			if (n_setPrime.count(cur + 1))
			{					
				i1 += i01;				
			}
			else
			{
				i0 += i01;
				i1 += i11;
			}
		}
		
		return { i0,i1 };
	}
	unordered_set<int> n_setPrime;
	long long m_llRet = 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
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
100
101
102
103
104
105
106
107

测试用例

template<class T,class T2>
void Assert(const T& t1, const T2& t2)
{
	assert(t1 == t2);
}

template<class T>
void Assert(const vector<T>& v1, const vector<T>& v2)
{
	if (v1.size() != v2.size())
	{
		assert(false);
		return;
	}
	for (int i = 0; i < v1.size(); i++)
	{
		Assert(v1[i], v2[i]);
	}

}

int main()
{		
	int n;
	vector<vector<int>> edges;
	{
		Solution sln;
		n = 5, edges = { {1,2},{1,3},{2,4},{2,5} };
		auto res = sln.countPaths(n, edges);
		Assert(res,4);
	}

	{
		Solution sln;
		n = 6, edges = { {1,2},{1,3},{2,4},{3,5},{3,6} };
		auto res = sln.countPaths(n, edges);
		Assert(res, 6);
	}
	
}

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

2023年9月

class Solution {
public:
long long countPaths(int n, vector& edges) {
std::set setPrime = { 2,3 };
for (int i = 4 ; i <= n; i++)
{
bool bPrime = true;
for (const int j : setPrime)
{
if (j * j > i)
{
break;
}
if (0 == i % j)
{
bPrime = false;
break;
}
}
if (bPrime)
{
setPrime.emplace(i);
}
}
CUnionFind uf(n);
for (const auto& v : edges)
{
if (setPrime.count(v[0]) || setPrime.count(v[1]))
{//联通所有非素数
continue;
}
uf.Union(v[0] - 1, v[1] - 1);
}
vector vNodeNum = uf.GetNodeCountOfRegion();
std::unordered_map mPreRegion;
for (const auto& v : edges)
{//素数和那些联通区域联通
const int n0 = v[0] - 1;
const int n1 = v[1] - 1;
if (setPrime.count(v[0]) && (!setPrime.count(v[1])))
{
mPreRegion[v[0]].emplace(uf.GetConnectRegionIndex(n1));
}
if (setPrime.count(v[1]) && (!setPrime.count(v[0])))
{
mPreRegion[v[1]].emplace(uf.GetConnectRegionIndex(n0));
}
}

	long long llRegion1 = 0, llRegion2 = 0;
	for (const auto& [tmp, setRegion] : mPreRegion)
	{
		int llSumNode = 0;
		for (const auto& region : setRegion)
		{
			llRegion1 += vNodeNum[region];//起点一个素数一个非素数
			llSumNode += vNodeNum[region];
		}
		for (const auto& region : setRegion)
		{
			llRegion2 += vNodeNum[region] * (llSumNode-vNodeNum[region]);
		}
	}
	return llRegion1 + llRegion2/2;
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

};

扩展阅读

视频课程

有效学习：明确的目标 及时的反馈 拉伸区（难度合适），可以先学简单的课程，请移步CSDN学院，听白银讲师（也就是鄙人）的讲解。
https://edu.csdn.net/course/detail/38771

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

相关

下载

想高屋建瓴的学习算法，请下载《喜缺全书算法册》doc版
https://download.csdn.net/download/he_zhidan/88348653

测试环境

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