LeetCode每日一题【c++版】

20240227-困难2867-统计树中的合法路径数目

一、概述

        由于比较难打，直接截图贴过来了！

二、思路

1. 预处理得到 [1,n]中的所有质数，其中 prime[i]是否为质数。
2. 根据二维整数整数构建图 g，其中 g[i]表示节点i的所有邻居节点。如果一条边的两个节点都不是质数，那么我们就将这两个节点合并到同一个连通分量中。
3. 在 [1,n]的范围内枚举所有质数i，考虑包含i的所有路径。
4. 由于i已经是质数，如果i是路径的一个端点，那么我们只需要累计与节点i相邻的所有连通分量的大小即可。如果i是路径上的某个中间点，那么我们需要累计相邻的任意两个连通分量的大小之积

需要知识点：并查集（此方法有点难，周末有时间会专门写一下并查集的专题）+ 枚举以及之前的一个预处理模块。

三、复杂度分析

时间复杂度：O(n*α(n))；其中n为节点数，α为阿克曼函数的反函数

空间复杂度：O(n)

四、代码

#include <iostream>
#include <vector>
#include <numeric>
using namespace std;

const int mx = 1e5 + 10;
bool prime[mx + 1];
int init = []() {
    for (int i = 2; i <= mx; ++i)
        prime[i] = true;
    for (int i = 2; i <= mx; ++i) {
        if (prime[i]) {
            for (int j = i + i; j <= mx; j += i) {
                prime[j] = false;
            }
        }
    }
    return 0;
}();

class UnionFind {
public:
    UnionFind(int n)
    {
        p = vector<int>(n);
        size = vector<int>(n, 1);
        iota(p.begin(), p.end(), 0);
    }

    bool unite(int a, int b)
    {
        int pa = find(a), pb = find(b);
        if (pa == pb) {
            return false;
        }
        if (size[pa] > size[pb]) {
            p[pb] = pa;
            size[pa] += size[pb];
        } else {
            p[pa] = pb;
            size[pb] += size[pa];
        }
        return true;
    }

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

    int getSize(int x)
    {
        return size[find(x)];
    }

private:
    vector<int> p, size;
};

class Solution {
public:
    long long countPaths(int n, vector<vector<int>> &edges)
    {
        vector<int> g[n + 1];
        UnionFind uf(n + 1);
        for (auto &e : edges) {
            int u = e[0], v = e[1];
            g[u].push_back(v);
            g[v].push_back(u);
            if (!prime[u] && !prime[v]) {
                uf.unite(u, v);
            }
        }
        long long ans = 0;
        for (int i = 1; i <= n; ++i) {
            if (prime[i]) {
                long long t = 0;
                for (int j : g[i]) {
                    if (!prime[j]) {
                        long long cnt = uf.getSize(j);
                        ans += cnt;
                        ans += cnt * t;
                        t += cnt;
                    }
                }
            }
        }
        return ans;
    }
};
int main()
{
    int n = 6;
    vector<vector<int>> edges = {{1, 2}, {1, 3}, {2, 4}, {3, 5}, {3, 6}};
    Solution solution;
    long long ans = solution.countPaths(n, edges);
    cout << ans << endl;

    n = 5;
    edges = {{1, 2}, {1, 3}, {2, 4}, {2, 5}};
    ans = solution.countPaths(n, edges);
    cout << ans << endl;
}