【树】【图论】【树路径】【深度优先搜索】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^5^
  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++**实现。