标签： hard

目标

给定一个由 n 个节点组成的网络，用 n x n 个邻接矩阵 graph 表示。在节点网络中，只有当 graph[i][j] = 1 时，节点 i 能够直接连接到另一个节点 j。

一些节点 initial 最初被恶意软件感染。只要两个节点直接连接，且其中至少一个节点受到恶意软件的感染，那么两个节点都将被恶意软件感染。这种恶意软件的传播将继续，直到没有更多的节点可以被这种方式感染。

假设 M(initial) 是在恶意软件停止传播之后，整个网络中感染恶意软件的最终节点数。

我们可以从 initial 中删除一个节点，并完全移除该节点以及从该节点到任何其他节点的任何连接。

请返回移除后能够使 M(initial) 最小化的节点。如果有多个节点满足条件，返回索引 最小的节点 。

示例 1：

输入：graph = [[1,1,0],[1,1,0],[0,0,1]], initial = [0,1]
输出：0

示例 2：

输入：graph = [[1,1,0],[1,1,1],[0,1,1]], initial = [0,1]
输出：1

示例 3：

输入：graph = [[1,1,0,0],[1,1,1,0],[0,1,1,1],[0,0,1,1]], initial = [0,1]
输出：1

说明：

• n == graph.length
• n == graph[i].length
• 2 <= n <= 300
• graph[i][j] 是 0 或 1.
• graph[i][j] == graph[j][i]
• graph[i][i] == 1
• 1 <= initial.length < n
• 0 <= initial[i] <= n - 1
• initial 中每个整数都不同

思路

这个和昨天的题的区别是移除节点之后原来连通的区域可能就断开了。刚开始想，昨天的需要排除掉同一连通区域存在多个感染节点的情况，今天这个就不能排除了。但是其影响的节点数也不能通过连通区域节点个数来计算。处理起来就比较复杂了，不能简单地根据直接相连的节点数来判断。当连通区域中含有多个感染节点时，需要区分边缘感染节点与中间感染节点，边缘感染节点又与孤立的感染节点相同，都是减少1。然后还要考虑连通区域仅有一个感染节点的情况。

区分 单一边缘感染节点 与 孤立感染节点

2、3 感染，返回3

0 - 1
|
3

2

无需区分 多个边缘感染节点 与 孤立感染节点

1、2、3 感染，返回1

0 - 1
|
3

2

区分 中间感染节点 与 孤立感染节点，并且不能仅根据直接相连的非感染节点来判断

0 、2、4、8 感染，返回8

     7
     | \
0    4 - 6
     |
1 -  8 - 3
     |
2    5

错了好多次，终于调试过了。正面确实不太好求解。总结一下就是：

1. 连通区域存在多个感染节点
   • 去掉边缘的感染节点，感染节点总数减少1，全是边缘感染节点与包含中间感染节点是一样的
   • 去掉非边缘感染节点，需要dfs获取不含感染节点路径的节点总数
2. 连通区域仅有1个感染节点（可以是孤立感染节点、边缘节点、中间节点）
   • 感染节点总数减少连通区域节点个数

最终答案需要获取以上减少最多的节点，如果存在多个，返回下标最小的。

代码里是按边缘感染节点与中间感染节点分的：

1. 边缘感染节点
   • 孤立感染节点，减1
   • 连通区域内有多个边缘感染节点，减1
   • 连通区域内仅有一个边缘感染节点，减连通区域节点个数
2. 中间感染节点（如果存在中间节点就不考虑边缘节点了，因为题目中限制了1 <= initial.length < n，一定存在可以减少2个的中间节点，分析到这里时我以为我发现了官网的漏洞，错误的实现也能通过，想要贡献测试用例呢，结果提示测试用例非有效值。如果是小于等于n这个解法就得多一个判断条件initial.length == n，直接取最小下标）
   • 仅有一个中间感染节点，连通区域节点个数
   • 有多个中间感染节点，dfs搜索不含感染节点路径上的非感染节点个数，如果有感染节点，那么它也是减1，不过这里不再比较了，原因上面也说过了。

代码中d存的是中间节点（包括指向自身的边大于2），如果d为空则表示连通区域全是边缘感染节点（边为2），或孤立感染节点（边为1）。

对于全是边缘感染节点与孤立感染节点的情况，取下标最小即可。而对于中间感染节点，通过dfs来查找连通个数。如果通过dfs查找的个数为1，并且它还是中间感染节点，那么它周围全是感染节点。按道理来说，应该与边缘节点一起取最小下标。但是题目给出了限制，那么一定存在一个可以减少2的中间节点。

通过上面的分析只是说明了该问题正向分析的复杂性，如果不是不断尝试，很难直接把上面的所有情况想清楚。所以，上面的分析也没有太大的用处，过一段时间重做这个题，还是会踩坑。

官网题解使用的是逆向思维，统计的是从每个非感染节点出发不经过感染节点所经历的个数，在dfs过程中使用状态机来标识感染节点的个数。如果只遇到了1个感染节点，那么累加刚才遍历的节点个数，而如果有多个，那么就只能减少它自己。因此，如果存在只遇到一个感染节点的情况，就取个数最大的。否则取下标最小的。

其实，只遇到一个感染节点的情况包括了上面的单一边缘感染节点、中间单一感染节点以及多个中间感染节点（dfs非感染个数不为0的情况，即路径上不含有感染节点）的情况，而遇到多个感染节点，则说明被多个感染节点包围/半包围（对应全是边缘节点、边缘与中间、全中间，后面两种情况上面的算法忽略掉了），并且取最小下标直接包括了孤立感染节点。

可以发现同样是一步处理，我们赋予它不同的内涵，其所应对的场景就大为不同。

代码

/**
 * @date 2024-04-17 8:46
 */
public class MinMalwareSpread928 {
    public int[] u;
    TreeSet<Integer> s;
    HashSet<Integer> d = new HashSet<>();
    List<Integer>[] g;

    public void merge(int x, int y) {
        HashSet<Integer> tmp = new HashSet<>();
        int rx = find(x, tmp);
        int ry = find(y, tmp);
        d.addAll(tmp);
        if (s.contains(rx) && s.contains(ry)) {
            if (rx > ry) {
                u[rx] = ry;
            } else if (rx < ry) {
                u[ry] = rx;
            }
        } else if (s.contains(ry)) {
            u[rx] = ry;
        } else {
            u[ry] = rx;
        }
    }

    public int find(int x, HashSet<Integer> tmp) {
        if (x != u[x]) {
            if (s.contains(x) && s.contains(u[x])) {
                if (g[x].size() > 2) {
                    tmp.add(x);
                }
                if (g[u[x]].size() > 2) {
                    tmp.add(u[x]);
                }
            }
            x = find(u[x], tmp);
        }
        return u[x];
    }

    public int find(int x) {
        if (x != u[x]) {
            x = find(u[x]);
        }
        return u[x];
    }

    public int count(int x) {
        int cnt = 0;
        int rt = find(x);
        for (int i = 0; i < u.length; i++) {
            if (rt == find(i)) {
                cnt++;
            }
        }
        return cnt;
    }

    public int countMalware(int x) {
        int cnt = 0;
        int rt = find(x);
        for (int i = 0; i < u.length; i++) {
            if (rt == find(i) && s.contains(i)) {
                cnt++;
            }
        }
        return cnt;
    }

    public int adjacencyUninfected(int x, int parent) {
        int cnt = 1;
        boolean[] visited = new boolean[u.length];
        for (Integer node : g[x]) {
            if (parent == node || node == x || visited[node]) {
                continue;
            }
            visited = new boolean[u.length];
            if (!s.contains(node)) {
                int subCnt = dfs(node, x, visited);
                if (subCnt != 0) {
                    cnt += subCnt;
                }
            }
        }
        return cnt;
    }

    public int dfs(int x, int parent, boolean[] visited) {
        if (s.contains(x)) {
            return 0;
        }
        int cnt = 1;
        for (Integer node : g[x]) {
            if (parent == node || node == x || visited[node]) {
                visited[node] = true;
                continue;
            }
            visited[node] = true;
            if (s.contains(node)) {
                return 0;
            }
            int subCnt = dfs(node, x, visited);
            if (subCnt == 0) {
                return 0;
            } else {
                cnt += subCnt;
            }
        }
        return cnt;
    }

    public int minMalwareSpread(int[][] graph, int[] initial) {
        int n = graph.length;
        g = new ArrayList[n];
        u = new int[n];
        for (int i = 0; i < n; i++) {
            g[i] = new ArrayList<>(n);
            u[i] = i;
        }

        s = new TreeSet<>();
        for (int i : initial) {
            s.add(i);
        }
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < n; j++) {
                if (graph[i][j] == 1) {
                    g[i].add(j);
                    merge(i, j);
                }
            }
        }
        int res = Integer.MAX_VALUE;
        int tmp = Integer.MAX_VALUE;
        TreeSet<Integer> ini = new TreeSet<>((x, y) -> count(y) - count(x) == 0 ? (adjacencyUninfected(y, -1) - adjacencyUninfected(x, -1) == 0 ? x - y : adjacencyUninfected(y, -1) - adjacencyUninfected(x, -1)) : count(y) - count(x));
        if (d.isEmpty()) {
            // d为空表示连通区域仅有1个感染节点
            for (int i : initial) {
                if (countMalware(i) == 1 && count(i) > 1) {
                    // 连通区域节点数大于1
                    if (tmp == Integer.MAX_VALUE) {
                        tmp = i;
                    } else {
                        int ci = count(i);
                        int ct = count(tmp);
                        if (ci > ct) {
                            // 取连通区域节点数大的
                            tmp = i;
                        } else if (ci == ct) {
                            // 如果相等取下标小的
                            tmp = Math.min(i, tmp);
                        }
                    }
                } else {
                    // 对于孤立节点，直接取索引最小的即可
                    res = Math.min(i, res);
                }
            }
            // 如果全部是孤立节点，取res，否则取tmp
            return tmp == Integer.MAX_VALUE ? res : tmp;
        } else {
            ini.addAll(d);
        }

        return ini.first();
    }
}

性能

目标

给出了一个由 n 个节点组成的网络，用 n × n 个邻接矩阵图 graph 表示。在节点网络中，当 graph[i][j] = 1 时，表示节点 i 能够直接连接到另一个节点 j。

一些节点 initial 最初被恶意软件感染。只要两个节点直接连接，且其中至少一个节点受到恶意软件的感染，那么两个节点都将被恶意软件感染。这种恶意软件的传播将继续，直到没有更多的节点可以被这种方式感染。

假设 M(initial) 是在恶意软件停止传播之后，整个网络中感染恶意软件的最终节点数。

如果从 initial 中移除某一节点能够最小化 M(initial)， 返回该节点。如果有多个节点满足条件，就返回索引最小的节点。

请注意，如果某个节点已从受感染节点的列表 initial 中删除，它以后仍有可能因恶意软件传播而受到感染。

示例 1：

输入：graph = [[1,1,0],[1,1,0],[0,0,1]], initial = [0,1]
输出：0

示例 2：

输入：graph = [[1,0,0],[0,1,0],[0,0,1]], initial = [0,2]
输出：0

示例 3：

输入：graph = [[1,1,1],[1,1,1],[1,1,1]], initial = [1,2]
输出：1

说明：

• n == graph.length
• n == graph[i].length
• 2 <= n <= 300
• graph[i][j] == 0 或 1.
• graph[i][j] == graph[j][i]
• graph[i][i] == 1
• 1 <= initial.length <= n
• 0 <= initial[i] <= n - 1
• initial 中所有整数均不重复

思路

从 初始 已感染恶意软件的节点集合中去掉一个节点使得整个网络的感染节点数量最小，返回这个节点。注意，从初始被感染的集合中去除，并不代表后续不会再被感染。如果还有与它连通的恶意节点，那么仍会被感染，最终计算感染节点时要算上。

因此，如果被感染节点是连通的，去掉任一感染节点后，总的感染节点数量不会改变。这时需要将索引最小的节点返回。

刚开始的想法是先排除相互的连通的感染节点，然后取剩余节点中连接节点个数最多的那个。

这个想法没错，但是具体实现的时候，仅仅判断直接相连的两个节点是否同时在感染列表显然是不对的，因为存在间接连接的情况。并且直接从感染集合移除还好影响后续其它节点的判断。

于是想到了使用并查集。

官网的解法类似，将连通的节点染成同一颜色，然后在感染节点中看是否有颜色唯一的节点，即该连通区域中只有一个感染节点，然后找出连通区域节点数最大的，如果有多个颜色唯一节点，返回下标最小的。如果没有颜色唯一的节点，那么移除任一感染节点，总的感染数都不会减少，直接取下标最小的即可。

判断区域是否连通可以使用并查集，也可以使用深度优先搜索。

代码

/**
 * @date 2024-04-16 8:29
 */
public class MinMalwareSpread924 {

    public int[] u;
    TreeSet<Integer> s;
    HashSet<Integer> d = new HashSet<>();

    public void merge(int x, int y) {
        HashSet<Integer> tmp = new HashSet<>();
        int rx = find(x, tmp);
        int ry = find(y, tmp);
        d.addAll(tmp);
        if (s.contains(rx) && s.contains(ry)) {
            if (rx > ry) {
                u[rx] = ry;
            } else if (rx < ry) {
                u[ry] = rx;
            }
        } else if (s.contains(ry)) {
            u[rx] = ry;
        } else {
            u[ry] = rx;
        }
    }

    public int find(int x, HashSet<Integer> tmp) {
        if (x != u[x]) {
            if (s.contains(x) && s.contains(u[x])) {
                tmp.add(x);
                tmp.add(u[x]);
            }
            x = find(u[x], tmp);
        }
        return u[x];
    }

    public int find(int x) {
        if (x != u[x]) {
            x = find(u[x]);
        }
        return u[x];
    }

    public int count(int x) {
        int cnt = 0;
        int rt = find(x);
        for (int i = 0; i < u.length; i++) {
            if (rt == find(i)) {
                cnt++;
            }
        }
        return cnt;
    }

    public int minMalwareSpread(int[][] graph, int[] initial) {
        int n = graph.length;
        List<Integer>[] g = new ArrayList[n];
        u = new int[n];
        for (int i = 0; i < n; i++) {
            g[i] = new ArrayList<>(n);
            u[i] = i;
        }

        s = new TreeSet<>();
        for (int i : initial) {
            s.add(i);
        }
        int res = s.first();
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < n; j++) {
                if (graph[i][j] == 1) {
                    g[i].add(j);
                    merge(i, j);
                }
            }
        }
        if (s.size() == d.size()) {
            return res;
        }
        TreeSet<Integer> ini = new TreeSet<>((x, y) -> count(y) - count(x) == 0 ? x - y : count(y) - count(x));
        for (int i : initial) {
            if (!d.contains(i)) {
                ini.add(i);
            }
        }

        return ini.first();
    }

}

性能

目标

给你一棵树，树上有 n 个节点，按从 0 到 n-1 编号。树以父节点数组的形式给出，其中 parent[i] 是节点 i 的父节点。树的根节点是编号为 0 的节点。

树节点的第 k 个祖先节点是从该节点到根节点路径上的第 k 个节点。

实现 TreeAncestor 类：

• TreeAncestor（int n， int[] parent） 对树和父数组中的节点数初始化对象。
• getKthAncestor(int node, int k) 返回节点 node 的第 k 个祖先节点。如果不存在这样的祖先节点，返回 -1 。

说明：

• 1 <= k <= n <= 5 * 10^4
• parent[0] == -1 表示编号为 0 的节点是根节点。
• 对于所有的 0 < i < n ，0 <= parent[i] < n 总成立
• 0 <= node < n
• 至多查询 5 * 10^4 次

思路

这个题让我们维护一个数据结构，来查找树中任意节点的第k个祖先节点。直接的想法是保存每一个节点的父节点，需要的时候直接根据下标获取。刚开始用的 int[][] 超出了空间限制，后来改成 List<Integer>[] 虽然多通过了几个测试用例，但是后面会超时。仔细分析最坏的情况下（所有节点仅有一个子树的情况），需要添加 n(n+1)/2 个父节点(首项为1，公差为1的等差数列求和)，时间复杂度是O(n^2)。

一个解决办法是不要保存重复的父节点，以只有一个子树的情况举例，最后一个节点第k个祖先，就是其父节点的第k-1个祖先。如果这个节点已经保存有祖先节点的信息，就无需重复计算了。

所以我的解决方案就是使用缓存，如果父节点的祖先信息没有保存，就将当前节点的祖先信息写入缓存，直到遇到存在缓存的祖先节点，如果它记录的祖先节点个数大于k - cnt就直接返回，否则继续向该缓存的祖先节点集合添加，直到遇到下一个有缓存的节点或者cnt == k。

这种方法虽然能够通过，但是与测试用例的的顺序是有关的，如果是从子节点逐步向前测试的话，缓存一直不命中，时间复杂度还是O(n^2)。

官方的解法使用的是倍增的思想，好像还挺常用的，算是个模板算法。核心思想是保存当前节点的父节点，爷爷节点，爷爷的爷爷节点......，即每个节点 x 的第 2^i 个祖先节点。这样不论k取什么值，都可以分解为不同的2的幂之和，然后向前查找即可。预处理的时间复杂度是O(nlogn)，查询的时间复杂度是O(logk)。

代码

/**
 * @date 2024-04-06 9:45
 */
public class TreeAncestor1483 {

    /**倍增的写法 */
    public static class TreeAncestor_v4 {

        int[][] dp;

        public TreeAncestor_v4(int n, int[] parent) {
            dp = new int[16][];
            dp[0] = parent;
            for (int i = 1; i < 16; i++) {
                dp[i] = new int[n];
                Arrays.fill(dp[i], -1);
            }

            for (int i = 1; i < 16; i++) {
                for (int j = 0; j < n; j++) {
                    if (dp[i - 1][j] != -1) {
                        dp[i][j] = dp[i - 1][dp[i - 1][j]];
                    }
                }
            }
        }

        public int getKthAncestor(int node, int k) {
            int p = node;
            int b = 0;
            int mod;
            while (k != 0) {
                mod = k & 1;
                if (mod == 1) {
                    p = dp[b][p];
                    if (p == -1) {
                        return -1;
                    }
                }
                k = k >> 1;
                b++;
            }
            return p;
        }
    }

    int[] parent;
    List<Integer>[] cache;

    public TreeAncestor1483(int n, int[] parent) {
        this.parent = parent;
        cache = new ArrayList[n];
        for (int i = 0; i < cache.length; i++) {
            cache[i] = new ArrayList<>();
        }
    }

    public int getKthAncestor(int node, int k) {
        if (node == -1) {
            return -1;
        }
        int cnt = 0;
        int p = node;
        while (cnt != k && p != -1) {
            if (cache[p].size() == 0) {
                cache[node].add(parent[p]);
                p = parent[p];
                cnt++;
            } else {
                if (cache[p].size() >= k - cnt) {
                    return cache[p].get(k - cnt - 1);
                } else {
                    cnt += cache[p].size();
                    node = p;
                    p = cache[p].get(cache[p].size() - 1);
                }
            }

        }
        return p;
    }
}

性能

这里是使用缓存写法的耗时，官方题解的耗时差不多也是这个样。

使用倍增的写法

目标

给你一个整数数组 nums （下标从 0 开始）和一个整数 k 。

一个子数组 (i, j) 的 分数 定义为 min(nums[i], nums[i+1], ..., nums[j]) * (j - i + 1) 。一个 好 子数组的两个端点下标需要满足 i <= k <= j 。

请你返回 好 子数组的最大可能 分数 。

示例 1：

输入：nums = [1,4,3,7,4,5], k = 3
输出：15
解释：最优子数组的左右端点下标是 (1, 5) ，分数为 min(4,3,7,4,5) * (5-1+1) = 3 * 5 = 15 。

示例 2：

输入：nums = [5,5,4,5,4,1,1,1], k = 0
输出：20
解释：最优子数组的左右端点下标是 (0, 4) ，分数为 min(5,5,4,5,4) * (4-0+1) = 4 * 5 = 20 。

提示：

• 1 <= nums.length <= 10^5
• 1 <= nums[i] <= 2 * 10^4
• 0 <= k < nums.length

思路

题目中定义的好子数组必须要包含下标k，且其元素最小值乘以它的长度应最大。相同长度的子数组其最小值通常不同，应取最小值中最大的，这样才能在窗口固定的情况下求得最大分数。

刚开始我把这个问题作为一个动态规划问题来求解：有一个窗口，在下标k的位置有一个固定轴，窗口可以左右滑动，拉伸，但窗口边缘不能越过k。然后求解窗口大小固定时，滑动窗口内最小元素取最大时的状态。接着扩展窗口，新窗口的取值依赖于上一窗口，只需在上一窗口的基础上左右各扩展一个元素进行比较即可。

但是我马上就遇到了问题，因为k的位置是不确定的，窗口左右滑动总会有一边先到达边界，然后怎么处理？上一个窗口取得较大的最小值可能是在k左侧，当窗口到达左侧边界后就无法再移动了，这样势必会有一部分覆盖到k右侧，我们无法再用一侧的最优解来掩盖另一侧了。而右边新加入窗口的元素与上一个状态选择的最小值无法确定新的最小值。因为窗口记录的是左右两侧的最优解，单独某一侧的状态并没有被记录。比如 nums=[10,9,8,7,6,5,3,2,2,4,9,4,11,3,4]，k=5，当窗口大小为6时，左侧的最小值是5，右侧最小值是2（但是我们并没有记录），我们记录的是较大的5。当窗口大小为7时，左侧窗口最小值为3（必须跨过k了），右侧新加入窗口的值是4，如果与上一个状态比较，我们可能会选择4，但是右侧最小值是2，我们应该选3。

于是我想可能需要分别记录左右两侧的状态。我们为什么要记录状态？上面记录状态是为了与新进入窗口的元素比较来选择最优解，我们现在记录左右两侧的什么呢？

随着思考的深入，我觉得应该放弃所谓滑动窗口这个概念了，不应该在左右两侧同时求解。

思考这个问题，窗口增大之后，其中元素的最小值会怎么变？反正最小值一定不会变大。于是只要新加入的元素比窗口内已经选定的最小值大就可以一直扩张，因为最小值没有变化，窗口大小越大分数就越大。当遇到比当前窗口内最小值小的元素时就需要比较窗口另一侧的值，哪边的更大就从哪边扩张。如此反复即可。

代码

/**
 * @date 2024-03-19 0:16
 */
public class MaximumScore {

    public int maximumScore_v2(int[] nums, int k) {
        if (nums.length == 1) {
            return nums[0];
        }
        int res = 0;
        int l = k - 1, r = k + 1;
        int lmin = nums[k], rmin = nums[k];
        while (l >= 0 || r < nums.length) {
            if (l >= 0) {
                lmin = Math.min(lmin, nums[l]);
            }
            if (r < nums.length) {
                rmin = Math.min(rmin, nums[r]);
            }
            if ((lmin >= rmin && l >= 0) || r >= nums.length) {
                l--;
                while (l >= 0 && lmin <= nums[l]) {
                    l--;
                }
                // r-l是窗口大小（不包括r），由于l多减了1，所以这里要减去
                res = Math.max(res, lmin * (r - l - 1));
            } else {
                r++;
                while (r < nums.length && rmin <= nums[r]) {
                    r++;
                }
                // r-l是窗口大小（不包括l）由于r多加了1，所以这里要减去
                res = Math.max(res, rmin * (r - l - 1));
            }
        }
        return res;
    }
}

性能

目标

给你一棵 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的自然数（非负整数）中，除了1和它本身以外不再有其他因数的自然数。

现在有一颗 n 个节点的无向树，要求任意两个连通节点间恰好有一个质数节点的路径数。树是一种无环连通图，问题可以转化为从树中选取两个节点，节点之间的路径只经过一个质数节点。由于没有环，所以两点之间的路径是唯一的。

1. 两个都是质数节点要排除掉。
2. 以质数节点为中心，与它邻接的非质数节点符合条件。即以质数节点为中心加上与之相连的非质数节点任取两个均可。我们可以称直接与质数节点相连的非质数节点为直接节点。
3. 直接节点连通的非质数节点也可能满足条件，需要要减去直接节点向外连通的路径，即直接节点加上其向外连通的节点之间任取两个的路径数。

按照上面的思路，先要找到所有的质数节点，涉及到质数判断。同时保存与之直接相连的非质数节点。然后保存非质数节点的边，使用Map保存，边的两个端点都保存进去，方便后续向外查找连通的节点。

得到满足条件的节点总数，根据排列组合公式C(n,2) = n!/(2!(n-2)!) = (n-1)n/2 求得路径总数D。

将外围节点k向外连通节点总数记为Ik，无效路径数为(Ik-1)Ik/2。

最终的结果就是D - Σ(Ik-1)Ik/2

代码

/**
 * @date 2024-02-27 0:22
 */
public class CountPaths {

    public Map<Integer, Set<Integer>> primeEdges = new HashMap<>();
    public Map<Integer, Set<Integer>> notPrimeEdges = new HashMap<>();
    public Map<Integer, Integer> indirectNodesNumMap = new HashMap<>();
    Set<Integer> counter = new HashSet<>();

    public long countPaths(int n, int[][] edges) {
        for (int i = 0; i < n - 1; i++) {
            int[] edge = edges[i];
            boolean i0 = isPrimeNumber(edge[0]);
            boolean i1 = isPrimeNumber(edge[1]);
            if (i0 && !i1) {
                primeEdges.computeIfAbsent(edge[0], k -> new HashSet<>());
                primeEdges.get(edge[0]).add(edge[1]);
            } else if (!i0 && i1) {
                primeEdges.computeIfAbsent(edge[1], k -> new HashSet<>());
                primeEdges.get(edge[1]).add(edge[0]);
            } else if(!i0){
                notPrimeEdges.computeIfAbsent(edge[0], k -> new HashSet<>());
                notPrimeEdges.computeIfAbsent(edge[1], k -> new HashSet<>());
                notPrimeEdges.get(edge[0]).add(edge[1]);
                notPrimeEdges.get(edge[1]).add(edge[0]);
            }
        }
        long res = 0;
        for (Integer primeNode : primeEdges.keySet()) {
            Set<Integer> nonPrimeNodesOfPrimeEdge = primeEdges.get(primeNode);
            counter.clear();
            int total = 0;
            for (int nonPrimeNode : nonPrimeNodesOfPrimeEdge) {
                counter.add(nonPrimeNode);
                if (indirectNodesNumMap.get(nonPrimeNode) == null) {
                    indirectNodesNumMap.put(nonPrimeNode, 1);
                    countEdges(nonPrimeNode, nonPrimeNode);
                }
                total += indirectNodesNumMap.get(nonPrimeNode);
            }
            total = total + 1;
            res += total * (total - 1L) / 2L;
            for (int nonPrimeNode : nonPrimeNodesOfPrimeEdge) {
                int indirectNodesNum = indirectNodesNumMap.get(nonPrimeNode);
                res -= indirectNodesNum * (indirectNodesNum - 1L) / 2L;
            }
        }
        return res;
    }

    public Set<Integer> countEdges(int key, int nonPrimeNode) {
        if (notPrimeEdges.get(nonPrimeNode) != null) {
            for (Integer node : notPrimeEdges.get(nonPrimeNode)) {
                if (!counter.contains(node)) {
                    indirectNodesNumMap.put(key, indirectNodesNumMap.get(key) + 1);
                    counter.add(node);
                    countEdges(key, node);
                }
            }
        }
        return counter;
    }

    public boolean isPrimeNumber(int num) {
        if (num == 1) {
            return false;
        }
        if (num == 2) {
            return true;
        }
        if (num % 2 == 0) {
            return false;
        }
        for (int i = 3;  i * i <= num; i+=2) {
            if (num % i == 0) {
                return false;
            }
        }
        return true;
    }

    public static void main(String[] args) {
//        int[][] edges = new int[][]{new int[]{1, 2}, new int[]{1, 3}, new int[]{2, 4}, new int[]{2, 5}};
        int[][] edges = new int[][]{new int[]{1, 2}, new int[]{4, 1}, new int[]{3, 4}};
        CountPaths main = new CountPaths();
//        System.out.println(main.countPaths(5, edges));
        System.out.println(main.countPaths(4, edges));
    }
}

性能

勉强通过。有时间再回来看看题解吧。