力扣 每日一题 统计树中的合法路径数目

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

文章目录

• 思路
• 复杂度
• Code

思路

‍ 灵神题解

代码实现时，为避免反复 DFS 同一个非质数连通块，可以把每个非质数所处的连通块的大小记录下来（类似记忆化搜索）。如果之前计算过，就无需再次 DFS。

复杂度

时间复杂度: $O(n)$

空间复杂度: $O(n)$

Code

class Solution {private final static int MX = (int) 1e5;private final static boolean[] np = new boolean[MX + 1]; // 质数=false 非质数=truestatic {np[1] = true;for (int i = 2; i * i <= MX; i++) {if (!np[i]) {for (int j = i * i; j <= MX; j += i) {np[j] = true;}}}}public long countPaths(int n, int[][] edges) {List[] g = new ArrayList[n + 1];Arrays.setAll(g, e -> new ArrayList());for (var e : edges) {int x = e[0], y = e[1];g[x].add(y);g[y].add(x);}long ans = 0;int[] size = new int[n + 1];var nodes = new ArrayList();for (int x = 1; x <= n; x++) {if (np[x]) { // 跳过非质数continue;}int sum = 0;for (int y : g[x]) { // 质数 x 把这棵树分成了若干个连通块if (!np[y]) {continue;}if (size[y] == 0) { // 尚未计算过nodes.clear();dfs(y, -1, g, nodes); // 遍历 y 所在连通块，在不经过质数的前提下，统计有多少个非质数for (int z : nodes) {size[z] = nodes.size();}}// 这 size[y] 个非质数与之前遍历到的 sum 个非质数，两两之间的路径只包含质数 xans += (long) size[y] * sum;sum += size[y];}ans += sum; // 从 x 出发的路径}return ans;}private void dfs(int x, int fa, List[] g, List nodes) {nodes.add(x);for (int y : g[x]) {if (y != fa && np[y]) {dfs(y, x, g, nodes);}}}}