Talentd
For employersSign in
JobsInternshipsPracticeResume toolsCompaniesArticles
For employers
FleetCode
Make and follow sheets.Sign in

Count Valid Paths in a Tree solution and explanation

HardDepth-First SearchDynamic ProgrammingMathNumber Theory8 min read36.2% acceptanceUpdated 9 Apr 2026
Asked atGameskraft
Practise Count Valid Paths in a Tree
28:32
2867. Count Valid Paths in a Tree | Weekly Leetcode 364codingMohan · 28:32 · 3K viewsWatch on YouTube
7 more videos on this problem
  1. 2867. Count Valid Paths in a Tree | Weekly Leetcode 364codingMohan · 28:32 · 3K views
  2. 2867. Count Valid Paths in a Tree || Leetcode Weekly Contest 364Ashish Kumar · 20:39 · 1K views
  3. 2867. Count Valid Paths in a Tree (Leetcode Hard)Programming Live with Larry · 1:10:51 · 593 views
  4. 【每日一题】LeetCode 2867. Count Valid Paths in a TreeHuifeng Guan · 19:58 · 513 views
  5. Leetcode 2867. Count Valid Paths in a TreeElite Code · 46:36 · 491 views
  6. 100047. Count Valid Paths in a Tree : LeetCode Weekly Contest 364 || SolutionBitByBit - Learn Coding · 0:26 · 197 views
  7. Disjoint Set Union | Leetcode 2867 | Why | Beginner friendlyCodeWithMeGuys · 13:47 · 127 views
  8. Count Valid Paths in Tree || LeetCode Weekly Contest 364 || Sieve, CombinatoricsMLE · 26:46 · 358 views

The problem

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.

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 <= 105
  • edges.length == n - 1
  • edges[i].length == 2
  • 1 <= ui, vi <= n
  • The input is generated such that edges represent a valid tree.

Approach overview

Problem Overview: You are given a tree with n nodes labeled from 1..n. A path is valid if it contains exactly one prime-numbered node. The task is to count all such paths. Because the structure is a tree, every pair of nodes has exactly one simple path, which lets you reason about paths through prime nodes.

Approach 1: DFS Traversal with Path Prime Check (O(n log log n) time, O(n) space)

Start by precomputing all prime numbers up to n using the Sieve of Eratosthenes. The key observation: every valid path must include exactly one prime node, so treat each prime node as the “center” of potential paths. For each prime node, run DFS into its neighboring components but only traverse non-prime nodes. Each DFS returns the size of a connected component of non-prime nodes attached to that prime.

Suppose the component sizes are s1, s2, s3.... Any pair of nodes chosen from two different components forms a valid path that passes through the prime node. You also count single-ended paths from the prime to each non-prime node. Accumulate combinations using a running prefix sum to avoid quadratic merging. The algorithm performs a single traversal per component and uses adjacency lists for the tree. This approach relies heavily on Depth-First Search and basic Number Theory for prime detection.

Approach 2: Dynamic Programming on Tree (Tree DP) (O(n log log n) time, O(n) space)

This version computes path counts during a single DFS using dynamic programming on the tree. For every node, maintain two states: paths in its subtree that contain 0 primes and paths that contain 1 prime. When merging a child subtree into the parent, combine states carefully so the resulting path still contains exactly one prime.

If the current node is non-prime, zero-prime paths extend normally, and one-prime paths can be formed by attaching a subtree that already contains one prime. If the current node is prime, any extension immediately introduces the single allowed prime, so you only combine with zero-prime child paths. Each merge step contributes new valid paths to the global answer.

This technique avoids explicitly iterating over components for each prime node. Instead, it pushes information upward during DFS. The method is typical Dynamic Programming on a Tree, where each node summarizes information from its children before returning to its parent.

Recommended for interviews: The DFS component-count approach is the most common solution. It shows that you recognized the “exactly one prime” constraint and converted the problem into counting combinations of non-prime components around a prime node. The Tree DP approach is equally optimal but more complex to reason about under time pressure. Demonstrating the DFS solution first and discussing the DP variant shows strong algorithmic depth.

Approach 1: DFS Traversal with Path Prime Check

This approach uses Depth First Search (DFS) to traverse the tree and count valid paths. For each node, we check paths from the current node to all other nodes using DFS. We store whether a path contains a prime number. If exactly one node in the path is prime, it is a valid path. However, this method alone would be too slow for large trees, so optimizations around path sharing and careful counting are necessary.

from collections import defaultdict

def sieve_of_eratosthenes(n):
    is_prime = [True] * (n + 1)
    p = 2
    while p * p <= n:
        if is_prime[p]:
            for i in range(p * p, n + 1, p):
                is_prime[i] = False
        p += 1
    is_prime[0], is_prime[1] = False, False
    return is_prime


def dfs(node, parent, graph, primes, visited):
    visited[node] = True
    path_count = 0
    for neighbor in graph[node]:
        if neighbor == parent:
            continue
        if not visited[neighbor]:
            if primes[neighbor] != primes[node]:
                path_count += 1
            path_count += dfs(neighbor, node, graph, primes, visited)
    return path_count


def count_valid_paths(n, edges):
    graph = defaultdict(list)
    for u, v in edges:
        graph[u].append(v)
        graph[v].append(u)
    primes = sieve_of_eratosthenes(n)
    visited = [False] * (n + 1)
    return dfs(1, -1, graph, primes, visited)

# Example Usage
n = 5
edges = [[1, 2], [1, 3], [2, 4], [2, 5]]
print(count_valid_paths(n, edges))
import java.util.*;

public class TreePaths {
    static List<Integer>[] graph;
    static boolean[] primes;
    static boolean[] visited;

    public static void sieveOfEratosthenes(int n) {
        Arrays.fill(primes, true);
        primes[0] = primes[1] = false;
        for (int p = 2; p * p <= n; p++) {
            if (primes[p]) {
                for (int i = p * p; i <= n; i += p) primes[i] = false;
            }
        }
    }

    public static int dfs(int node, int parent) {
        visited[node] = true;
        int pathCount = 0;
        for (int neighbor : graph[node]) {
            if (neighbor == parent) continue;
            if (!visited[neighbor]) {
                if (primes[neighbor] != primes[node]) pathCount++;
                pathCount += dfs(neighbor, node);
            }
        }
        return pathCount;
    }

    public static int countValidPaths(int n, int[][] edges) {
        graph = new ArrayList[n + 1];
        for (int i = 0; i <= n; i++) graph[i] = new ArrayList<>();
        primes = new boolean[n + 1];
        visited = new boolean[n + 1];
        sieveOfEratosthenes(n);

        for (int[] edge : edges) {
            int u = edge[0], v = edge[1];
            graph[u].add(v);
            graph[v].add(u);
        }
        return dfs(1, -1);
    }

    public static void main(String[] args) {
        int n = 5;
        int[][] edges = {{1, 2}, {1, 3}, {2, 4}, {2, 5}};
        System.out.println(countValidPaths(n, edges));
    }
}
#include <iostream>
#include <vector>
#include <cmath>
using namespace std;

vector<int> graph[100001];
bool primes[100001];
bool visited[100001];

void sieve_of_eratosthenes(int n) {
    fill(primes, primes + n + 1, true);
    primes[0] = primes[1] = false;
    for (int p = 2; p * p <= n; p++) {
        if (primes[p]) {
            for (int i = p * p; i <= n; i += p) primes[i] = false;
        }
    }
}

int dfs(int node, int parent) {
    visited[node] = true;
    int path_count = 0;
    for (int neighbor : graph[node]) {
        if (neighbor == parent) continue;
        if (!visited[neighbor]) {
            if (primes[neighbor] != primes[node]) path_count++;
            path_count += dfs(neighbor, node);
        }
    }
    return path_count;
}

int count_valid_paths(int n, vector<vector<int>>& edges) {
    sieve_of_eratosthenes(n);
    for (const auto& edge : edges) {
        int u = edge[0], v = edge[1];
        graph[u].push_back(v);
        graph[v].push_back(u);
    }
    fill(visited, visited + n + 1, false);
    return dfs(1, -1);
}

int main() {
    int n = 5;
    vector<vector<int>> edges = {{1, 2}, {1, 3}, {2, 4}, {2, 5}};
    cout << count_valid_paths(n, edges) << endl;
    return 0;
}
function sieveOfEratosthenes(n) {
    const isPrime = new Array(n + 1).fill(true);
    isPrime[0] = isPrime[1] = false;
    for (let p = 2; p * p <= n; p++) {
        if (isPrime[p]) {
            for (let i = p * p; i <= n; i += p) isPrime[i] = false;
        }
    }
    return isPrime;
}

function dfs(node, parent, graph, primes, visited) {
    visited[node] = true;
    let pathCount = 0;
    for (let neighbor of graph[node]) {
        if (neighbor === parent) continue;
        if (!visited[neighbor]) {
            if (primes[neighbor] !== primes[node]) pathCount++;
            pathCount += dfs(neighbor, node, graph, primes, visited);
        }
    }
    return pathCount;
}

function countValidPaths(n, edges) {
    const graph = Array.from({ length: n + 1 }, () => []);
    for (const [u, v] of edges) {
        graph[u].push(v);
        graph[v].push(u);
    }
    const primes = sieveOfEratosthenes(n);
    const visited = new Array(n + 1).fill(false);
    return dfs(1, -1, graph, primes, visited);
}

// Example Usage
const n = 5;
const edges = [[1, 2], [1, 3], [2, 4], [2, 5]];
console.log(countValidPaths(n, edges));
#include <stdio.h>
#include <stdbool.h>
#include <stdlib.h>

#define MAX_NODES 100000

int graph[MAX_NODES + 1][MAX_NODES + 1];
bool primes[MAX_NODES + 1];
bool visited[MAX_NODES + 1];

void sieve_of_eratosthenes(int n) {
    for (int i = 2; i <= n; i++) primes[i] = true;
    for (int p = 2; p * p <= n; p++) {
        if (primes[p]) {
            for (int i = p * p; i <= n; i += p) primes[i] = false;
        }
    }
}

int dfs(int node, int n) {
    visited[node] = true;
    int count = 0;
    for (int i = 1; i <= n; i++) {
        if (graph[node][i] && !visited[i]) {
            if (primes[i] != primes[node]) count++;
            count += dfs(i, n);
        }
    }
    return count;
}

int count_valid_paths(int n, int edges[][2]) {
    memset(graph, 0, sizeof(graph));
    memset(visited, false, sizeof(visited));
    sieve_of_eratosthenes(n);

    for (int i = 0; i < n - 1; i++) {
        int u = edges[i][0];
        int v = edges[i][1];
        graph[u][v] = 1;
        graph[v][u] = 1;
    }

    return dfs(1, n);
}

int main() {
    int n = 5;
    int edges[4][2] = {{1, 2}, {1, 3}, {2, 4}, {2, 5}};
    printf("%d\n", count_valid_paths(n, edges));
    return 0;
}
using System;
using System.Collections.Generic;

public class TreePaths {
    static List<int>[] graph;
    static bool[] primes;
    static bool[] visited;

    public static void SieveOfEratosthenes(int n) {
        Array.Fill(primes, true);
        primes[0] = primes[1] = false;
        for (int p = 2; p * p <= n; p++) {
            if (primes[p]) {
                for (int i = p * p; i <= n; i += p) primes[i] = false;
            }
        }
    }

    public static int DFS(int node, int parent) {
        visited[node] = true;
        int pathCount = 0;
        foreach (int neighbor in graph[node]) {
            if (neighbor == parent) continue;
            if (!visited[neighbor]) {
                if (primes[neighbor] != primes[node]) pathCount++;
                pathCount += DFS(neighbor, node);
            }
        }
        return pathCount;
    }

    public static int CountValidPaths(int n, int[][] edges) {
        graph = new List<int>[n + 1];
        for (int i = 0; i <= n; i++) graph[i] = new List<int>();
        primes = new bool[n + 1];
        visited = new bool[n + 1];
        SieveOfEratosthenes(n);

        foreach (int[] edge in edges) {
            int u = edge[0], v = edge[1];
            graph[u].Add(v);
            graph[v].Add(u);
        }
        return DFS(1, -1);
    }

    public static void Main(string[] args) {
        int n = 5;
        int[][] edges = new int[][] {
            new int[] {1, 2},
            new int[] {1, 3},
            new int[] {2, 4},
            new int[] {2, 5}
        };
        Console.WriteLine(CountValidPaths(n, edges));
    }
}

The Python solution uses defaultdict for the adjacency list representation of the graph. Primes are checked with a Sieve of Eratosthenes, and DFS is used to traverse the tree and count valid paths.

Complexity

Time Complexity: O(n), because the problem involves traversing the entire tree.
Space Complexity: O(n) for storing the graph and visited flags.

Try this approach in the editor

Approach 2: Dynamic Programming on Tree (Tree DP)

This optimized approach uses Tree Dynamic Programming to calculate valid paths more efficiently. By employing tree-centric calculations, instead of evaluating every path individually, the algorithm calculates path properties at each node and aggregates results efficiently using dynamic programming techniques.

from collections import defaultdict

def sieve_of_eratosthenes(n):
    is_prime = [True] * (n + 1)
    p = 2
    while p * p <= n:
        if is_prime[p]:
            for i in range(p * p, n + 1, p):
                is_prime[i] = False
        p += 1
    is_prime[0], is_prime[1] = False, False
    return is_prime


def dfs(node, parent, graph, primes, subtreePrimes):
    for neighbor in graph[node]:
        if neighbor == parent:
            continue
        if primes[neighbor]:
            subtreePrimes[node] += 1
        dfs(neighbor, node, graph, primes, subtreePrimes)


def count_valid_paths(n, edges):
    graph = defaultdict(list)
    for u, v in edges:
        graph[u].append(v)
        graph[v].append(u)
    primes = sieve_of_eratosthenes(n)
    subtreePrimes = [0] * (n + 1)
    dfs(1, -1, graph, primes, subtreePrimes)

    validPaths = 0
    for u in range(1, n + 1):
        for v in graph[u]:
            if u < v:
                if primes[u] != primes[v]:
                    validPaths += 1

    return validPaths

# Example Usage
n = 5
edges = [[1, 2], [1, 3], [2, 4], [2, 5]]
print(count_valid_paths(n, edges))
import java.util.*;

public class TreePaths {
    static List<Integer>[] graph;
    static boolean[] primes;
    static boolean[] visited;
    static int[] subtreePrimes;

    public static void sieveOfEratosthenes(int n) {
        Arrays.fill(primes, true);
        primes[0] = primes[1] = false;
        for (int p = 2; p * p <= n; p++) {
            if (primes[p]) {
                for (int i = p * p; i <= n; i += p) primes[i] = false;
            }
        }
    }

    public static int dfs(int node, int parent) {
        visited[node] = true;
        int children = 0;
        for (int neighbor : graph[node]) {
            if (neighbor != parent && !visited[neighbor]) {
                children++;
                if (primes[neighbor]) subtreePrimes[node]++;
                dfs(neighbor, node);
            }
        }
        return children;
    }

    public static int countValidPaths(int n, int[][] edges) {
        graph = new ArrayList[n + 1];
        for (int i = 0; i <= n; i++) graph[i] = new ArrayList<>();
        primes = new boolean[n + 1];
        visited = new boolean[n + 1];
        subtreePrimes = new int[n + 1];
        sieveOfEratosthenes(n);

        for (int[] edge : edges) {
            int u = edge[0], v = edge[1];
            graph[u].add(v);
            graph[v].add(u);
        }
        Arrays.fill(subtreePrimes, 0);
        Arrays.fill(visited, false);
        dfs(1, -1);

        int validPaths = 0;
        for (int i = 1; i <= n; i++) {
            for (int neighbor : graph[i]) {
                if (i < neighbor) {
                    if (primes[i] != primes[neighbor]) validPaths++;
                }
            }
        }

        return validPaths;
    }

    public static void main(String[] args) {
        int n = 5;
        int[][] edges = {{1, 2}, {1, 3}, {2, 4}, {2, 5}};
        System.out.println(countValidPaths(n, edges));
    }
}
#include <iostream>
#include <vector>
#include <cstring>
using namespace std;

vector<int> graph[100001];
bool primes[100001];
bool visited[100001];
int subtreePrimes[100001];

void sieve_of_eratosthenes(int n) {
    fill(primes, primes + n + 1, true);
    primes[0] = primes[1] = false;
    for (int p = 2; p * p <= n; p++) {
        if (primes[p]) {
            for (int i = p * p; i <= n; i += p) primes[i] = false;
        }
    }
}

int dfs(int node, int parent) {
    visited[node] = true;
    int children = 0;
    for (int neighbor : graph[node]) {
        if (neighbor != parent && !visited[neighbor]) {
            children++;
            if (primes[neighbor]) subtreePrimes[node]++;
            dfs(neighbor, node);
        }
    }
    return children;
}

int countValidPaths(int n, vector<vector<int>>& edges) {
    sieve_of_eratosthenes(n);
    for (const auto& edge : edges) {
        int u = edge[0], v = edge[1];
        graph[u].push_back(v);
        graph[v].push_back(u);
    }
    memset(subtreePrimes, 0, sizeof(subtreePrimes));
    memset(visited, false, sizeof(visited));
    dfs(1, -1);

    int validPaths = 0;
    for (int i = 1; i <= n; i++) {
        for (int neighbor : graph[i]) {
            if (i < neighbor) {
                if (primes[i] != primes[neighbor]) validPaths++;
            }
        }
    }

    return validPaths;
}

int main() {
    int n = 5;
    vector<vector<int>> edges = {{1, 2}, {1, 3}, {2, 4}, {2, 5}};
    cout << countValidPaths(n, edges) << endl;
    return 0;
}
function sieveOfEratosthenes(n) {
    const isPrime = new Array(n + 1).fill(true);
    isPrime[0] = isPrime[1] = false;
    for (let p = 2; p * p <= n; p++) {
        if (isPrime[p]) {
            for (let i = p * p; i <= n; i += p) isPrime[i] = false;
        }
    }
    return isPrime;
}

function dfs(node, parent, graph, primes, subtreePrimes) {
    for (let neighbor of graph[node]) {
        if (neighbor === parent) continue;
        if (primes[neighbor]) subtreePrimes[node]++;
        dfs(neighbor, node, graph, primes, subtreePrimes);
    }
}

function countValidPaths(n, edges) {
    const graph = Array.from({ length: n + 1 }, () => []);
    for (const [u, v] of edges) {
        graph[u].push(v);
        graph[v].push(u);
    }
    const primes = sieveOfEratosthenes(n);
    const subtreePrimes = new Array(n + 1).fill(0);
    dfs(1, -1, graph, primes, subtreePrimes);

    let validPaths = 0;
    for (let u = 1; u <= n; u++) {
        for (let v of graph[u]) {
            if (u < v) {
                if (primes[u] !== primes[v]) {
                    validPaths++;
                }
            }
        }
    }

    return validPaths;
}

// Example Usage
const n = 5;
const edges = [[1, 2], [1, 3], [2, 4], [2, 5]];
console.log(countValidPaths(n, edges));
#include <stdio.h>
#include <stdbool.h>
#include <string.h>
#include <stdlib.h>

#define MAX_NODES 100000

int graph[MAX_NODES + 1][MAX_NODES + 1];
bool primes[MAX_NODES + 1];
bool visited[MAX_NODES + 1];
int subtreePrimes[MAX_NODES + 1];

void sieve_of_eratosthenes(int n) {
    for (int i = 2; i <= n; i++) primes[i] = true;
    for (int p = 2; p * p <= n; p++) {
        if (primes[p]) {
            for (int i = p * p; i <= n; i += p) primes[i] = false;
        }
    }
}

int dfs(int node, int n) {
    visited[node] = true;
    int count = 0;
    for (int i = 1; i <= n; i++) {
        if (graph[node][i] && !visited[i]) {
            if (primes[i]) subtreePrimes[node]++;
            count += dfs(i, n);
        }
    }
    return count;
}

int count_valid_paths(int n, int edges[][2]) {
    memset(graph, 0, sizeof(graph));
    memset(visited, false, sizeof(visited));
    memset(subtreePrimes, 0, sizeof(subtreePrimes));
    sieve_of_eratosthenes(n);

    for (int i = 0; i < n - 1; i++) {
        int u = edges[i][0];
        int v = edges[i][1];
        graph[u][v] = 1;
        graph[v][u] = 1;
    }

    dfs(1, n);

    int validPaths = 0;
    for (int i = 1; i <= n; i++) {
        for (int j = 1; j <= n; j++) {
            if (graph[i][j] && i < j) {
                if (primes[i] != primes[j]) validPaths++;
            }
        }
    }

    return validPaths;
}

int main() {
    int n = 5;
    int edges[4][2] = {{1, 2}, {1, 3}, {2, 4}, {2, 5}};
    printf("%d\n", count_valid_paths(n, edges));
    return 0;
}
using System;
using System.Collections.Generic;

public class TreePaths {
    static List<int>[] graph;
    static bool[] primes;
    static bool[] visited;
    static int[] subtreePrimes;

    public static void SieveOfEratosthenes(int n) {
        Array.Fill(primes, true);
        primes[0] = primes[1] = false;
        for (int p = 2; p * p <= n; p++) {
            if (primes[p]) {
                for (int i = p * p; i <= n; i += p) primes[i] = false;
            }
        }
    }

    public static void DFS(int node, int parent) {
        foreach (int neighbor in graph[node]) {
            if (neighbor != parent) {
                if (primes[neighbor]) subtreePrimes[node]++;
                DFS(neighbor, node);
            }
        }
    }

    public static int CountValidPaths(int n, int[][] edges) {
        graph = new List<int>[n + 1];
        for (int i = 0; i <= n; i++) graph[i] = new List<int>();
        primes = new bool[n + 1];
        visited = new bool[n + 1];
        subtreePrimes = new int[n + 1];
        SieveOfEratosthenes(n);

        foreach (int[] edge in edges) {
            int u = edge[0], v = edge[1];
            graph[u].Add(v);
            graph[v].Add(u);
        }
        Array.Fill(subtreePrimes, 0);
        DFS(1, -1);

        int validPaths = 0;
        for (int i = 1; i <= n; i++) {
            foreach (int neighbor in graph[i]) {
                if (i < neighbor) {
                    if (primes[i] != primes[neighbor]) validPaths++;
                }
            }
        }

        return validPaths;
    }

    public static void Main(string[] args) {
        int n = 5;
        int[][] edges = new int[][] {
            new int[] {1, 2},
            new int[] {1, 3},
            new int[] {2, 4},
            new int[] {2, 5}
        };
        Console.WriteLine(CountValidPaths(n, edges));
    }
}

This Python version optimizes based on subtree prime calculations, reducing redundant path assessments. By analyzing adjacency connectivity, valid paths are determined efficiently.

Complexity

Time Complexity: O(n), addressing paths through neighbors.
Space Complexity: O(n) for graph and path storage.

Try this approach in the editor

Approach 3: Preprocessing + Union-Find + Enumeration

We can preprocess to get all the prime numbers in [1, n], where prime[i] indicates whether i is a prime number.

Next, we build a graph g based on the two-dimensional integer array, where g[i] represents all the neighbor nodes of node i. If both nodes of an edge are not prime numbers, we merge these two nodes into the same connected component.

Then, we enumerate all prime numbers i in the range of [1, n], considering all paths that include i.

Since i is already a prime number, if i is an endpoint of the path, we only need to accumulate the sizes of all connected components adjacent to node i. If i is a middle point on the path, we need to accumulate the product of the sizes of any two adjacent connected components.

The time complexity is O(n times \alpha(n)), and the space complexity is O(n). Here, n is the number of nodes, and \alpha is the inverse function of the Ackermann function.

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;
}
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
}
Try this approach in the editor

Complexity comparison

ApproachComplexity
DFS Traversal with Path Prime Check

Time Complexity: O(n), because the problem involves traversing the entire tree.
Space Complexity: O(n) for storing the graph and visited flags.

Dynamic Programming on Tree (Tree DP)

Time Complexity: O(n), addressing paths through neighbors.
Space Complexity: O(n) for graph and path storage.

Preprocessing + Union-Find + Enumeration–

Complexity analysis

ApproachTimeSpaceWhen to Use
DFS Traversal with Prime Component CountingO(n log log n)O(n)Best general solution. Clear reasoning: treat each prime node as a hub and count combinations of adjacent non-prime components.
Dynamic Programming on Tree (Tree DP)O(n log log n)O(n)Useful when solving multiple path-count variants in a tree or when you prefer single-pass DP aggregation.

Frequently asked questions

Is Count Valid Paths in a Tree easy or hard?

Count Valid Paths in a Tree is classified as Hard. The challenge is recognizing that valid paths must contain exactly one prime node and converting the problem into counting combinations of non-prime components around each prime node using DFS or tree dynamic programming.

Count Valid Paths in a Tree Python/Java solution

Typical implementations build the adjacency list, compute primes with the Sieve of Eratosthenes, then run DFS from prime nodes to measure sizes of connected non-prime components. The same logic works in Python, Java, C++, and other languages with O(n log log n) time complexity.

How to solve Count Valid Paths in a Tree in O(n)?

Strict O(n) is difficult because prime detection typically uses a sieve costing O(n log log n). After preprocessing primes, the counting step is linear. DFS is used to measure sizes of non-prime components connected to each prime node, and combinations of those sizes produce the number of valid paths.

What is the best approach for Count Valid Paths in a Tree?

The most practical approach treats every prime node as the center of valid paths. Run DFS into adjacent non-prime components, compute their sizes, and count pair combinations across components. With prime preprocessing using a sieve, the total complexity becomes O(n log log n) time and O(n) space.

Is Count Valid Paths in a Tree asked at Google/Amazon/Meta?

Hard tree problems involving DFS, prime preprocessing, and path counting frequently appear in interviews at companies like Google, Amazon, and Meta. Variants that combine graph traversal with number theory or combinatorial counting are especially common in senior-level interview rounds.

What data structure is used in Count Valid Paths in a Tree?

The core data structure is an adjacency list to represent the tree efficiently. The algorithm also uses DFS recursion or a stack for traversal, arrays for prime flags from the sieve, and counters to track component sizes when combining paths.

What is the time complexity of Count Valid Paths in a Tree?

The optimal solutions run in O(n log log n) time due to the Sieve of Eratosthenes used to precompute primes. The DFS traversals themselves are linear, O(n), because each edge in the tree is processed a constant number of times. Space complexity is O(n) for adjacency lists and visited tracking.

Comments

Next problem · similarCount Paths That Can Form a Palindrome in a TreeHard
Now write it yourselfClose this page and solve Count Valid Paths in a Tree in the editor, without looking.
Practise Count Valid Paths in a Tree