Count Valid Paths in a Tree
There is an undirected tree with n nodes labeled from 1 to n. You are given the integer n and a 2D integer array edges of length n - 1, where edges[i] = [ui, vi] indicates that there is an edge between nodes ui and vi in the tree.
Return the number of valid paths in the tree.
A path (a, b) is valid if there exists exactly one prime number among the node labels in the path from a to b.
Note that:
- The path
(a, b)is a sequence of distinct nodes starting with nodeaand ending with nodebsuch that every two adjacent nodes in the sequence share an edge in the tree. - Path
(a, b)and path(b, a)are considered the same and counted only once.
Example 1:
Input: n = 5, edges = [[1,2],[1,3],[2,4],[2,5]]
Output: 4
Explanation: The pairs with exactly one prime number on the path between them are:
- (1, 2) since the path from 1 to 2 contains prime number 2.
- (1, 3) since the path from 1 to 3 contains prime number 3.
- (1, 4) since the path from 1 to 4 contains prime number 2.
- (2, 4) since the path from 2 to 4 contains prime number 2.
It can be shown that there are only 4 valid paths.Example 2:
Input: n = 6, edges = [[1,2],[1,3],[2,4],[3,5],[3,6]]
Output: 6
Explanation: The pairs with exactly one prime number on the path between them are:
- (1, 2) since the path from 1 to 2 contains prime number 2.
- (1, 3) since the path from 1 to 3 contains prime number 3.
- (1, 4) since the path from 1 to 4 contains prime number 2.
- (1, 6) since the path from 1 to 6 contains prime number 3.
- (2, 4) since the path from 2 to 4 contains prime number 2.
- (3, 6) since the path from 3 to 6 contains prime number 3.
It can be shown that there are only 6 valid paths.
Constraints:
1 <= n <= 105edges.length == n - 1edges[i].length == 21 <= ui, vi <= n- The input is generated such that
edgesrepresent a valid tree.
解题思路
把简单路径按照经过的最高点进行分类，即根据路径两个端点的最近公共祖先来分类。对于以 $u$ 为根的子树，假设与其相邻的子节点为 $v_i$，由于经过 $u$ 的路径中必须恰好有一个质数节点，因此在子树 $v_i$ 中，含端点 $v_i$ 的路径中最多只能有一个质数节点。
对此定义 $f(u, 0/1)$ 表示在子树 $u$ 中含端点 $u$ 的路径中恰好有 $0/1$ 个质数节点的数量。当 $u$ 是质数节点时，有 $\displaylines{\begin{cases} f(u,0) = 0 \\ f(u,1) = 1 + \sum{f(v_i,0)} \end{cases}}$。当 $u$ 不是质数节点时，有 $\displaylines{\begin{cases} f(u,0) = 1 + \sum{f(v_i,0)} \\ f(u,1) = \sum{f(v_i,1)} \end{cases}}$。
再考虑两个端点的最近公共祖先是 $u$ 且经过恰好一个质数节点的简单路径的数量。如果 $u$ 是质数节点，则分成两个部分，一个是端点含 $u$ 的，对应的数量是 $\sum{f(v_i,0)}$。另一个两个端点来自不同的子树 $v_i$，对应的数量是 $\sum{f(v_i,0) \sum\limits_{j>i}{f(v_j,0)}}$。如果 $u$ 不是质数节点，也分成两个部分，一个是端点含 $u$ 的，对应的数量是 $\sum{f(v_i,1)}$。另一个两个端点来自不同的子树 $v_i$，对应的数量是 $\sum{f(v_i,1) \sum\limits_{j \ne i}{f(v_j,0)}}$。
AC 代码如下，时间复杂度为 $O(n)$：
class Solution {
public:
long long countPaths(int n, vector<vector<int>>& edges) {
vector<vector<int>> g(n + 1);
for (auto &p : edges) {
g[p[0]].push_back(p[1]);
g[p[1]].push_back(p[0]);
}
vector<int> primes;
vector<bool> vis(n + 1);
vis[1] = true;
for (int i = 2; i <= n; i++) {
if (!vis[i]) primes.push_back(i);
for (int j = 0; primes[j] * i <= n; j++) {
vis[primes[j] * i] = true;
if (i % primes[j] == 0) break;
}
}
vector<vector<int>> f(n + 1, vector<int>(2));
long long ans = 0;
function<void(int, int)> dfs = [&](int u, int p) {
int s0 = 0, s1 = 0;
for (auto &v : g[u]) {
if (v == p) continue;
dfs(v, u);
s0 += f[v][0];
s1 += f[v][1];
}
if (!vis[u]) {
f[u][0] = 0;
f[u][1] = s0 + 1;
ans += s0;
for (auto &v : g[u]) {
if (v == p) continue;
s0 -= f[v][0];
ans += 1ll * f[v][0] * s0;
}
}
else {
f[u][0] = s0 + 1;
f[u][1] = s1;
ans += s1;
for (auto &v : g[u]) {
if (v == p) continue;
ans += 1ll * f[v][1] * (s0 - f[v][0]);
}
}
};
dfs(1, -1);
return ans;
}
};