2867. 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 node a and ending with node b such 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.

class Solution {
    vector<bool> prime;
    vector<long long> uf, cnt;
    long long find(int u) {
        return uf[u] == u ? u : uf[u] = find(uf[u]);
    }
    void uni(int u, int v) {
        long long pu = find(u), pv = find(v);
        cnt[pu] = cnt[pv] = cnt[pu] + cnt[pv];
        uf[pu] = uf[pv] = min(pu,pv);
    }
public:
    long long countPaths(int n, vector<vector<int>>& edges) {
        prime = vector<bool>(n + 1);
        cnt = uf = vector<long long>(n + 1);
        for(int i = 1; i <= n; i++) cnt[i] = 1, uf[i] = i;
        prime[0] = prime[1] = true;
        for(long long i = 2; i <= n; i++) {
            if(prime[i]) continue;
            for(long long j = i * i; j <= n; j += i) prime[j] = true;
        }
        vector<pair<long long, long long>> exc;
        for(auto e : edges) {
            int u = e[0], v = e[1];
            if(!prime[u] or !prime[v]) {
                if(!prime[u] and !prime[v]) continue;
                exc.push_back({u,v});
            } else {
                uni(u,v);
            }
        }
        long long res = 0;
        unordered_map<long long, long long> freq;
        for(auto& [u,v] : exc) {
            u = find(u), v = find(v);
            if(!prime[u]) {
                if(!freq.count(u)) freq[u] = 1;
                res += freq[u] * cnt[v];
                freq[u] += cnt[v];
            }
            if(!prime[v]) {
                if(!freq.count(v)) freq[v] = 1;
                res += freq[v] * cnt[u];
                freq[v] += cnt[u];
            }

        }
        return res;
    }
};