Problem statement

An undirected tree has n nodes labelled 1 to n, joined by the n - 1 edges in edges.

A path between two different nodes a and b is valid if exactly one of the node labels on it (endpoints included) is a prime number. The path from a to b and the path from b to a count as the same path.

Return the number of valid paths.

CodeKairo bounds n at 10^4 so the count always fits in a 32-bit integer.

Example 1

Input: n = 5, edges = [[1,2],[1,3],[2,4],[2,5]]
Output: 4
Explanation: The valid paths are 1–2, 1–3, 1–2–4 and 2–4: each contains exactly one of the primes 2, 3, 5.

Example 2

Input: n = 6, edges = [[1,2],[1,3],[2,4],[3,5],[3,6]]
Output: 6
Explanation: 1–2, 1–3, 1–2–4, 2–4, 3–6 and 1–3–6.

Example 3

Input: n = 1, edges = []
Output: 0
Explanation: There is no pair of different nodes.

Constraints

• 1 <= n <= 10^4
• edges.length == n - 1
• edges[i].length == 2
• 1 <= u, v <= n
• edges represents a valid tree

Hints for Count Valid Paths in a Tree

1. Hint 1

   Remove every prime node: the non-prime nodes fall apart into connected groups. A valid path contains exactly one prime, so everything else on it lies inside such groups.

2. Hint 2

   Fix the prime p on the path. Each end of the path is either p itself or a node in one of the non-prime groups touching p — and the two ends must come from different groups.

3. Hint 3

   Find the groups with union-find over edges joining two non-primes. For a prime p whose neighbouring groups have sizes s1, s2, …, add s_i (paths ending at p) plus s_i · (s1 + … + s_{i-1}) (paths crossing p) for each i.

How to solve Count Valid Paths in a Tree

A valid path has a single prime p; the rest of the path consists of non-prime nodes on either side of p, each side inside one connected group of non-primes adjacent to p. Precompute the size of every non-prime group, and count, for each prime, the pairs of ends it can join.

Approach

1. Sieve the primes up to n.
2. Union-find: join u and v for every edge whose ends are both non-prime; group sizes come from the roots.
3. For each prime p, walk its non-prime neighbours; for a neighbour's group of size s, add s · (1 + seen) to the answer, then seen += s (seen starts at 0 for each prime).
4. Return the total.

Why it works

Each valid path corresponds to exactly one prime p and one or two ends in distinct groups adjacent to p (two neighbours of p can never share a group, since the tree has no cycle). s · 1 counts paths from p into the group, and s · seen counts paths between this group and the earlier ones through p; every such path is counted once.

Complexity

• Time — O(n α(n) + n log log n)
• Space — O(n)

Pitfalls

• 1 is not prime.
• Paths with zero primes or with two or more primes are both invalid — never count paths inside a single non-prime group.
• In the real problem n reaches 10^5 and the count needs 64 bits; with CodeKairo's bound of 10^4 it fits in 32.

Reference solution

Python

from typing import List

def countPaths(n: int, edges: List[List[int]]) -> int:
    is_prime = [False, False] + [True] * (n - 1)
    for i in range(2, int(n ** 0.5) + 1):
        if is_prime[i]:
            for j in range(i * i, n + 1, i):
                is_prime[j] = False
    parent = list(range(n + 1))
    size = [1] * (n + 1)

    def find(x):
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x

    adj = [[] for _ in range(n + 1)]
    for u, v in edges:
        adj[u].append(v)
        adj[v].append(u)
        if not is_prime[u] and not is_prime[v]:
            a, b = find(u), find(v)
            if a != b:
                if size[a] < size[b]:
                    a, b = b, a
                parent[b] = a
                size[a] += size[b]
    total = 0
    for p in range(2, n + 1):
        if not is_prime[p]:
            continue
        seen = 0
        for v in adj[p]:
            if not is_prime[v]:
                s = size[find(v)]
                total += s * (1 + seen)
                seen += s
    return total

Java

static int[] dsuParent;

static int dsuFind(int x) {
    while (dsuParent[x] != x) { dsuParent[x] = dsuParent[dsuParent[x]]; x = dsuParent[x]; }
    return x;
}

public static int countPaths(int n, int[][] edges) {
    boolean[] composite = new boolean[n + 1];
    composite[0] = true;
    if (n >= 1) composite[1] = true;
    for (int a = 2; (long) a * a <= n; a++) if (!composite[a]) for (int b = a * a; b <= n; b += a) composite[b] = true;
    dsuParent = new int[n + 1];
    int[] size = new int[n + 1];
    List<List<Integer>> adj = new ArrayList<>();
    for (int k = 0; k <= n; k++) { dsuParent[k] = k; size[k] = 1; adj.add(new ArrayList<>()); }
    for (int[] e : edges) {
        int u = e[0], v = e[1];
        adj.get(u).add(v); adj.get(v).add(u);
        if (composite[u] && composite[v]) {
            int ru = dsuFind(u), rv = dsuFind(v);
            if (ru != rv) {
                if (size[ru] < size[rv]) { int t = ru; ru = rv; rv = t; }
                dsuParent[rv] = ru;
                size[ru] += size[rv];
            }
        }
    }
    long total = 0;
    for (int p = 2; p <= n; p++) {
        if (composite[p]) continue;
        long seen = 0;
        for (int w : adj.get(p)) {
            if (!composite[w]) continue;
            long s = size[dsuFind(w)];
            total += s * (1 + seen);
            seen += s;
        }
    }
    return (int) total;
}

C++

int countPaths(int n, vector<vector<int>>& edges) {
    vector<bool> composite(n + 1, false);
    composite[0] = true;
    if (n >= 1) composite[1] = true;
    for (int a = 2; (long long) a * a <= n; a++) if (!composite[a]) for (int b = a * a; b <= n; b += a) composite[b] = true;
    vector<int> parent(n + 1), sz(n + 1, 1);
    iota(parent.begin(), parent.end(), 0);
    function<int(int)> find = [&](int x) {
        while (parent[x] != x) { parent[x] = parent[parent[x]]; x = parent[x]; }
        return x;
    };
    vector<vector<int>> adj(n + 1);
    for (auto& e : edges) {
        int u = e[0], v = e[1];
        adj[u].push_back(v); adj[v].push_back(u);
        if (composite[u] && composite[v]) {
            int ru = find(u), rv = find(v);
            if (ru != rv) {
                if (sz[ru] < sz[rv]) swap(ru, rv);
                parent[rv] = ru;
                sz[ru] += sz[rv];
            }
        }
    }
    long long total = 0;
    for (int p = 2; p <= n; p++) {
        if (composite[p]) continue;
        long long seen = 0;
        for (int w : adj[p]) {
            if (!composite[w]) continue;
            long long s = sz[find(w)];
            total += s * (1 + seen);
            seen += s;
        }
    }
    return (int) total;
}

JavaScript

var countPaths = function(n, edges) {
    var isPrime = new Array(n + 1);
    for (var i = 0; i <= n; i++) isPrime[i] = i >= 2;
    for (var a = 2; a * a <= n; a++) if (isPrime[a]) for (var b = a * a; b <= n; b += a) isPrime[b] = false;
    var parent = new Array(n + 1), size = new Array(n + 1), adj = [];
    for (var k = 0; k <= n; k++) { parent[k] = k; size[k] = 1; adj.push([]); }
    var find = function(x) {
        while (parent[x] !== x) { parent[x] = parent[parent[x]]; x = parent[x]; }
        return x;
    };
    for (var e = 0; e < edges.length; e++) {
        var u = edges[e][0], v = edges[e][1];
        adj[u].push(v); adj[v].push(u);
        if (!isPrime[u] && !isPrime[v]) {
            var ru = find(u), rv = find(v);
            if (ru !== rv) {
                if (size[ru] < size[rv]) { var t = ru; ru = rv; rv = t; }
                parent[rv] = ru;
                size[ru] += size[rv];
            }
        }
    }
    var total = 0;
    for (var p = 2; p <= n; p++) {
        if (!isPrime[p]) continue;
        var seen = 0;
        for (var j = 0; j < adj[p].length; j++) {
            var w = adj[p][j];
            if (isPrime[w]) continue;
            var s = size[find(w)];
            total += s * (1 + seen);
            seen += s;
        }
    }
    return total;
};

Also on the editorial tab: C, C#, Go, Kotlin, PHP, Ruby, Rust, Swift, TypeScript.

More dynamic programming problems

• Subtraction in Linked List Hard
• GCD Sort of an Array Hard
• Create Components With Same Value Hard
• Cat and Mouse Hard
• Find the Number of Subsequences With Equal GCD Hard
• Find the Count of Monotonic Pairs I Hard