LeetCode #2867 — HARD

Count Valid Paths in a Tree

Break down a hard problem into reliable checkpoints, edge-case handling, and complexity trade-offs.

The Problem

Problem Statement

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] = [u_i, v_i] indicates that there is an edge between nodes u_i and v_i 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.

Example 1:

Input: n = 5, edges = [[1,2],[1,3],[2,4],[2,5]]
Output: 4
Explanation: The pairs with exactly one prime number on the path between them are: 
- (1, 2) since the path from 1 to 2 contains prime number 2. 
- (1, 3) since the path from 1 to 3 contains prime number 3.
- (1, 4) since the path from 1 to 4 contains prime number 2.
- (2, 4) since the path from 2 to 4 contains prime number 2.
It can be shown that there are only 4 valid paths.

Example 2:

Input: n = 6, edges = [[1,2],[1,3],[2,4],[3,5],[3,6]]
Output: 6
Explanation: The pairs with exactly one prime number on the path between them are: 
- (1, 2) since the path from 1 to 2 contains prime number 2.
- (1, 3) since the path from 1 to 3 contains prime number 3.
- (1, 4) since the path from 1 to 4 contains prime number 2.
- (1, 6) since the path from 1 to 6 contains prime number 3.
- (2, 4) since the path from 2 to 4 contains prime number 2.
- (3, 6) since the path from 3 to 6 contains prime number 3.
It can be shown that there are only 6 valid paths.

Constraints:

1 <= n <= 10⁵
edges.length == n - 1
edges[i].length == 2
1 <= u_i, v_i <= n
The input is generated such that edges represent a valid tree.

Step 01

Brute Force Baseline

Problem summary: 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.

Baseline thinking

Start with the most direct exhaustive search. That gives a correctness anchor before optimizing.

Pattern signal: Math · Dynamic Programming · Tree

Example 1

5
[[1,2],[1,3],[2,4],[2,5]]

Example 2

6
[[1,2],[1,3],[2,4],[3,5],[3,6]]

Core Insight

What unlocks the optimal approach

Use the sieve of Eratosthenes to find all prime numbers in the range <code>[1, n]</code>.****
Root the tree at any node.
Let <code>dp[i][0] = the number of vertical paths starting from i containing no prime nodes </code>, and <code>dp[i][1] = the number of vertical paths starting from i containing one prime node </code>.
If <code>i</code> is not prime, <code>dp[i][0] = sum(dp[child][0]) + 1</code>, and <code>dp[i][1] = sum(dp[child][1])</code> for each <code>child</code> of <code>i</code> in the rooted tree.
If <code>i</code> is prime, <code>dp[i][0] = 0</code>, and <code>dp[i][1] = sum(dp[child][0]) + 1</code> for each <code>child</code> of <code>i</code> in the rooted tree.
For each node <code>i</code>, and using the computed <code>dp</code> matrix, count the number of unordered pairs <code>(a,b)</code> such that <code>lca(a,b) = i</code>, and there exists exactly one prime number on the path from <code>a</code> to <code>b</code>.

Interview move: turn each hint into an invariant you can check after every iteration/recursion step.

Step 03

Algorithm Walkthrough

Iteration Checklist

Define state (indices, window, stack, map, DP cell, or recursion frame).
Apply one transition step and update the invariant.
Record answer candidate when condition is met.
Continue until all input is consumed.

Use the first example testcase as your mental trace to verify each transition.

Step 04

Edge Cases

Minimum Input

Single element / shortest valid input

Validate boundary behavior before entering the main loop or recursion.

Duplicates & Repeats

Repeated values / repeated states

Decide whether duplicates should be merged, skipped, or counted explicitly.

Extreme Constraints

Largest constraint values

Re-check complexity target against constraints to avoid time-limit issues.

Invalid / Corner Shape

Empty collections, zeros, or disconnected structures

Handle special-case structure before the core algorithm path.

Step 05

Full Annotated Code

Source-backed implementations are provided below for direct study and interview prep.

// Accepted solution for LeetCode #2867: Count Valid Paths in a Tree
class PrimeTable {
    private final boolean[] prime;

    public PrimeTable(int n) {
        prime = new boolean[n + 1];
        Arrays.fill(prime, true);
        prime[0] = false;
        prime[1] = false;
        for (int i = 2; i <= n; ++i) {
            if (prime[i]) {
                for (int j = i + i; j <= n; j += i) {
                    prime[j] = false;
                }
            }
        }
    }

    public boolean isPrime(int x) {
        return prime[x];
    }
}

class UnionFind {
    private final int[] p;
    private final int[] size;

    public UnionFind(int n) {
        p = new int[n];
        size = new int[n];
        for (int i = 0; i < n; ++i) {
            p[i] = i;
            size[i] = 1;
        }
    }

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

    public boolean union(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;
    }

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

class Solution {
    private static final PrimeTable PT = new PrimeTable(100010);

    public long countPaths(int n, int[][] edges) {
        List<Integer>[] g = new List[n + 1];
        Arrays.setAll(g, i -> new ArrayList<>());
        UnionFind uf = new UnionFind(n + 1);
        for (int[] e : edges) {
            int u = e[0], v = e[1];
            g[u].add(v);
            g[v].add(u);
            if (!PT.isPrime(u) && !PT.isPrime(v)) {
                uf.union(u, v);
            }
        }
        long ans = 0;
        for (int i = 1; i <= n; ++i) {
            if (PT.isPrime(i)) {
                long t = 0;
                for (int j : g[i]) {
                    if (!PT.isPrime(j)) {
                        long cnt = uf.size(j);
                        ans += cnt;
                        ans += cnt * t;
                        t += cnt;
                    }
                }
            }
        }
        return ans;
    }
}

// Accepted solution for LeetCode #2867: Count Valid Paths in a Tree
const mx int = 1e5 + 10

var prime [mx]bool

func init() {
	for i := 2; i < mx; i++ {
		prime[i] = true
	}
	for i := 2; i < mx; i++ {
		if prime[i] {
			for j := i + i; j < mx; j += i {
				prime[j] = false
			}
		}
	}
}

type unionFind struct {
	p, size []int
}

func newUnionFind(n int) *unionFind {
	p := make([]int, n)
	size := make([]int, n)
	for i := range p {
		p[i] = i
		size[i] = 1
	}
	return &unionFind{p, size}
}

func (uf *unionFind) find(x int) int {
	if uf.p[x] != x {
		uf.p[x] = uf.find(uf.p[x])
	}
	return uf.p[x]
}

func (uf *unionFind) union(a, b int) bool {
	pa, pb := uf.find(a), uf.find(b)
	if pa == pb {
		return false
	}
	if uf.size[pa] > uf.size[pb] {
		uf.p[pb] = pa
		uf.size[pa] += uf.size[pb]
	} else {
		uf.p[pa] = pb
		uf.size[pb] += uf.size[pa]
	}
	return true
}

func (uf *unionFind) getSize(x int) int {
	return uf.size[uf.find(x)]
}

func countPaths(n int, edges [][]int) (ans int64) {
	uf := newUnionFind(n + 1)
	g := make([][]int, n+1)
	for _, e := range edges {
		u, v := e[0], e[1]
		g[u] = append(g[u], v)
		g[v] = append(g[v], u)
		if !prime[u] && !prime[v] {
			uf.union(u, v)
		}
	}
	for i := 1; i <= n; i++ {
		if prime[i] {
			t := 0
			for _, j := range g[i] {
				if !prime[j] {
					cnt := uf.getSize(j)
					ans += int64(cnt + cnt*t)
					t += cnt
				}
			}
		}
	}
	return
}

# Accepted solution for LeetCode #2867: Count Valid Paths in a Tree
class UnionFind:
    def __init__(self, n):
        self.p = list(range(n))
        self.size = [1] * n

    def find(self, x):
        if self.p[x] != x:
            self.p[x] = self.find(self.p[x])
        return self.p[x]

    def union(self, a, b):
        pa, pb = self.find(a), self.find(b)
        if pa == pb:
            return False
        if self.size[pa] > self.size[pb]:
            self.p[pb] = pa
            self.size[pa] += self.size[pb]
        else:
            self.p[pa] = pb
            self.size[pb] += self.size[pa]
        return True


mx = 10**5 + 10
prime = [True] * (mx + 1)
prime[0] = prime[1] = False
for i in range(2, mx + 1):
    if prime[i]:
        for j in range(i * i, mx + 1, i):
            prime[j] = False


class Solution:
    def countPaths(self, n: int, edges: List[List[int]]) -> int:
        g = [[] for _ in range(n + 1)]
        uf = UnionFind(n + 1)
        for u, v in edges:
            g[u].append(v)
            g[v].append(u)
            if prime[u] + prime[v] == 0:
                uf.union(u, v)

        ans = 0
        for i in range(1, n + 1):
            if prime[i]:
                t = 0
                for j in g[i]:
                    if not prime[j]:
                        cnt = uf.size[uf.find(j)]
                        ans += cnt
                        ans += t * cnt
                        t += cnt
        return ans

// Accepted solution for LeetCode #2867: Count Valid Paths in a Tree
/**
 * [2867] Count Valid Paths in a Tree
 */
pub struct Solution {}

// submission codes start here

impl Solution {
    fn dfs(
        graph: &Vec<Vec<usize>>,
        primes: &Vec<bool>,
        seen: &mut Vec<usize>,
        i: usize,
        pre: usize,
    ) {
        seen.push(i);

        for next in &graph[i] {
            let next = *next;
            if next != pre && !primes[next] {
                Solution::dfs(graph, primes, seen, next, i);
            }
        }
    }

    pub fn count_paths(n: i32, edges: Vec<Vec<i32>>) -> i64 {
        let n = n as usize;
        let mut primes = vec![true; n + 1];
        primes[1] = false;
        for i in 2..=n {
            if primes[i] {
                if i * i > n {
                    continue;
                }

                for j in (i * i..=n).step_by(i) {
                    primes[j] = false;
                }
            }
        }
        primes = dbg!(primes);

        let mut graph = vec![vec![]; n + 1];
        for edge in edges {
            let x = edge[0] as usize;
            let y = edge[1] as usize;

            graph[x].push(y);
            graph[y].push(x);
        }

        let mut count = vec![0; n + 1];
        let mut result = 0;

        for root in 1..=n {
            if !primes[root] {
                continue;
            }

            let mut current = 0;
            for i in &graph[root] {
                let i = *i;
                if (primes[i]) {
                    continue;
                }

                if count[i] == 0 {
                    let mut seen = Vec::new();
                    Solution::dfs(&graph, &primes, &mut seen, i, 0);
                    for k in &seen {
                        count[*k] = seen.len() as i64;
                    }
                }

                result += count[i] * current;
                current += count[i];
            }

            result += current;
        }

        result
    }
}

// submission codes end

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_2867() {
        assert_eq!(
            6,
            Solution::count_paths(
                6,
                vec![vec![1, 2], vec![1, 3], vec![2, 4], vec![3, 5], vec![3, 6]]
            )
        );
    }
}

// Accepted solution for LeetCode #2867: Count Valid Paths in a Tree
const mx = 100010;
const prime = Array(mx).fill(true);
prime[0] = prime[1] = false;
for (let i = 2; i <= mx; ++i) {
    if (prime[i]) {
        for (let j = i + i; j <= mx; j += i) {
            prime[j] = false;
        }
    }
}

class UnionFind {
    p: number[];
    size: number[];
    constructor(n: number) {
        this.p = Array(n)
            .fill(0)
            .map((_, i) => i);
        this.size = Array(n).fill(1);
    }

    find(x: number): number {
        if (this.p[x] !== x) {
            this.p[x] = this.find(this.p[x]);
        }
        return this.p[x];
    }

    union(a: number, b: number): boolean {
        const [pa, pb] = [this.find(a), this.find(b)];
        if (pa === pb) {
            return false;
        }
        if (this.size[pa] > this.size[pb]) {
            this.p[pb] = pa;
            this.size[pa] += this.size[pb];
        } else {
            this.p[pa] = pb;
            this.size[pb] += this.size[pa];
        }
        return true;
    }

    getSize(x: number): number {
        return this.size[this.find(x)];
    }
}

function countPaths(n: number, edges: number[][]): number {
    const uf = new UnionFind(n + 1);
    const g: number[][] = Array(n + 1)
        .fill(0)
        .map(() => []);
    for (const [u, v] of edges) {
        g[u].push(v);
        g[v].push(u);
        if (!prime[u] && !prime[v]) {
            uf.union(u, v);
        }
    }
    let ans = 0;
    for (let i = 1; i <= n; ++i) {
        if (prime[i]) {
            let t = 0;
            for (let j of g[i]) {
                if (!prime[j]) {
                    const cnt = uf.getSize(j);
                    ans += cnt + t * cnt;
                    t += cnt;
                }
            }
        }
    }
    return ans;
}

Step 06

Interactive Study Demo

Use this to step through a reusable interview workflow for this problem.

Custom checkpoints (one per line)

Press Step or Run All to begin.

Step 07

Complexity Analysis

Time

O(n × \alpha(n)

Space

O(n)

Approach Breakdown

RECURSIVE

O(2ⁿ) time

O(n) space

Pure recursion explores every possible choice at each step. With two choices per state (take or skip), the decision tree has 2ⁿ leaves. The recursion stack uses O(n) space. Many subproblems are recomputed exponentially many times.

DYNAMIC PROGRAMMING

O(n × m) time

O(n × m) space

Each cell in the DP table is computed exactly once from previously solved subproblems. The table dimensions determine both time and space. Look for the state variables — each unique combination of state values is one cell. Often a rolling array can reduce space by one dimension.

Shortcut: Count your DP state dimensions → that’s your time. Can you drop one? That’s your space optimization.

Coach Notes

Common Mistakes

Review these before coding to avoid predictable interview regressions.

Overflow in intermediate arithmetic

Wrong move: Temporary multiplications exceed integer bounds.

Usually fails on: Large inputs wrap around unexpectedly.

Fix: Use wider types, modular arithmetic, or rearranged operations.

State misses one required dimension

Wrong move: An incomplete state merges distinct subproblems and caches incorrect answers.

Usually fails on: Correctness breaks on cases that differ only in hidden state.

Fix: Define state so each unique subproblem maps to one DP cell.

Forgetting null/base-case handling

Wrong move: Recursive traversal assumes children always exist.

Usually fails on: Leaf nodes throw errors or create wrong depth/path values.

Fix: Handle null/base cases before recursive transitions.