作者推荐
【广度优先搜索】【网格】【割点】【 推荐】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 <= 105
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; }
};
