作者推荐
【广度优先搜索】【网格】【割点】【 推荐】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 形成一棵合法的树。
深度优先搜索
以任意节点为根，任意两个节点的路径一定是： 起点 → \rightarrow → 公共祖先 → \rightarrow → 终点，特例：起点或终点就是公共祖先。我们枚举公共祖先。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;
};
测试用例
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);
}
}
2023年9月
class Solution {
public:
long long countPaths(int n, vector<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<int, unordered_set> 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;
}
};
扩展阅读
视频课程
有效学习：明确的目标 及时的反馈 拉伸区（难度合适），可以先学简单的课程，请移步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++**实现。