Count Valid Paths in a Tree - Practice Coding Problems

Count Valid Paths in a Tree - Problem

Math Dynamic Programming Tree Depth-First Search Number Theory Hard

Count Valid Paths in a Tree

Imagine you're exploring a mystical tree where each node is labeled with a number from 1 to n. Your quest is to find all the magical paths - those that contain exactly one prime number along the way!

You're given:
• An integer n representing the total number of nodes
• A 2D array edges where each edge connects two nodes, forming a tree

Your mission: Count how many valid paths exist in this tree.

A path from node a to node b is considered valid if the sequence of node labels along this path contains exactly one prime number. Remember that paths (a,b) and (b,a) are the same path and should only be counted once!

Example: In a tree with nodes [1,2,3,4] where 2 and 3 are prime, a path containing nodes [1,2,4] would be valid (exactly one prime: 2), but a path [2,3,4] would not be valid (two primes: 2 and 3).

Input & Output

example_1.py — Small Tree

$ Input: n = 4, edges = [[1,2],[1,3],[2,4]]

› Output: 4

💡 Note: The valid paths are: (1,2), (1,3), (2,4), (3,4). Each contains exactly one prime number. Path (1,2) has prime 2, path (1,3) has prime 3, path (2,4) has prime 2, and path (3,4) passes through nodes 3→1→2→4 containing primes 3 and 2, so it's invalid. Actually, direct paths: (1,2) has prime 2, (1,3) has prime 3, (2,4) has prime 2, and (1,4) passes through 1→2→4 with prime 2.

example_2.py — Linear Tree

$ Input: n = 5, edges = [[1,2],[2,3],[3,4],[4,5]]

› Output: 6

💡 Note: In this linear tree 1-2-3-4-5, valid paths with exactly one prime are: (1,2), (1,3), (2,4), (2,5), (4,5), (3,4). Path (2,3) has two primes so it's invalid.

example_3.py — No Prime Tree

$ Input: n = 3, edges = [[1,4],[4,6]]

› Output: 0

💡 Note: None of the nodes (1,4,6) are prime numbers, so no path contains exactly one prime. All paths have zero primes.

Visualization

Tap to expand

Understanding the Visualization

Identify Prime Cities

Mark all cities with prime numbers as treasure cities (red nodes)

Build Connection Map

Create adjacency list showing which cities connect to which

Count Valid Routes

Use DFS to count paths passing through exactly one treasure city

Optimize with DP

Group paths by prime node relationships to avoid redundant calculations

Key Takeaway

🎯 Key Insight: Use DFS to traverse the tree while maintaining a prime counter, and leverage tree properties to avoid checking every possible path individually.

Time & Space Complexity

Time Complexity

⏱️

O(n³)

O(n²) pairs of nodes, each requiring O(n) path traversal in worst case

⚠ Quadratic Growth

Space Complexity

O(n)

Space for adjacency list, visited array, and recursion stack

⚡ Linearithmic Space

Constraints

1 ≤ n ≤ 10⁵
edges.length == n - 1
edges[i].length == 2
1 ≤ u_i, v_i ≤ n
The input represents a valid tree

Asked in

G Google 15 a Amazon 8 f Meta 12 ⊞ Microsoft 6

This problem requires counting paths in a tree with exactly one prime number. The optimal approach uses DFS with dynamic programming, treating each node as a potential path endpoint and efficiently counting valid paths by grouping nodes based on their connectivity to prime nodes. Key insight: leverage the tree structure where there's exactly one path between any two nodes.

Common Approaches

Approach	Time	Space	Notes
✓ Brute Force (Check All Paths)	O(n³)	O(n)	Generate all possible paths and count primes in each
DFS with Tree Properties (Optimized)	O(n²)	O(n)	Use DFS to count paths efficiently by leveraging tree structure

Brute Force (Check All Paths) — Algorithm Steps

Build adjacency list representation of the tree
Generate Sieve of Eratosthenes to identify prime numbers up to n
For each pair of nodes (i,j) where i < j, find the path between them using DFS
Count prime numbers in each path
Increment result if exactly one prime is found

Visualization

Tap to expand

Step-by-Step Walkthrough

Generate all node pairs

Create pairs (1,2), (1,3), (1,4), (2,3), (2,4), (3,4)...

Find path for each pair

Use DFS to find the unique path between each pair of nodes

Count primes in path

Check each node in the path to see if it's prime

Validate and count

If exactly one prime found, increment result counter

Code -

solution.c — C

#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <string.h>

typedef struct {
    int* adj;
    int size;
    int capacity;
} AdjList;

long long countPaths(int n, int** edges, int edgesSize, int* edgesColSize) {
    AdjList* graph = (AdjList*)malloc((n + 1) * sizeof(AdjList));
    for (int i = 0; i <= n; i++) {
        graph[i].adj = (int*)malloc(10 * sizeof(int));
        graph[i].size = 0;
        graph[i].capacity = 10;
    }
    
    for (int i = 0; i < edgesSize; i++) {
        int u = edges[i][0], v = edges[i][1];
        graph[u].adj[graph[u].size++] = v;
        graph[v].adj[graph[v].size++] = u;
    }
    
    bool* isPrime = (bool*)malloc((n + 1) * sizeof(bool));
    for (int i = 0; i <= n; i++) isPrime[i] = true;
    isPrime[0] = isPrime[1] = false;
    
    for (int i = 2; i * i <= n; i++) {
        if (isPrime[i]) {
            for (int j = i * i; j <= n; j += i) {
                isPrime[j] = false;
            }
        }
    }
    
    long long result = 0;
    bool* visited = (bool*)malloc((n + 1) * sizeof(bool));
    int* path = (int*)malloc(n * sizeof(int));
    
    for (int i = 1; i <= n; i++) {
        for (int j = i + 1; j <= n; j++) {
            memset(visited, false, (n + 1) * sizeof(bool));
            int pathLen = findPath(i, j, graph, visited, path, 0);
            if (pathLen > 0) {
                int primeCount = 0;
                for (int k = 0; k < pathLen; k++) {
                    if (isPrime[path[k]]) primeCount++;
                }
                if (primeCount == 1) result++;
            }
        }
    }
    
    free(visited);
    free(path);
    free(isPrime);
    for (int i = 0; i <= n; i++) free(graph[i].adj);
    free(graph);
    
    return result;
}

int findPath(int start, int end, AdjList* graph, bool* visited, int* path, int pathLen) {
    path[pathLen] = start;
    if (start == end) return pathLen + 1;
    
    visited[start] = true;
    for (int i = 0; i < graph[start].size; i++) {
        int neighbor = graph[start].adj[i];
        if (!visited[neighbor]) {
            int len = findPath(neighbor, end, graph, visited, path, pathLen + 1);
            if (len > 0) {
                visited[start] = false;
                return len;
            }
        }
    }
    visited[start] = false;
    return 0;
}

Time & Space Complexity

Time Complexity

⏱️

O(n³)

O(n²) pairs of nodes, each requiring O(n) path traversal in worst case

⚠ Quadratic Growth

Space Complexity

O(n)

Space for adjacency list, visited array, and recursion stack

⚡ Linearithmic Space

Constraints

1 ≤ n ≤ 10⁵
edges.length == n - 1
edges[i].length == 2
1 ≤ u_i, v_i ≤ n
The input represents a valid tree

23.5K Views

Medium Frequency

~35 min Avg. Time

847 Likes

Ln 1, Col 1

Smart Actions

💡 Explanation

AI Ready

💡 Suggestion Tab to accept Esc to dismiss

// Output will appear here after running code

Code Editor Closed

Click the red button to reopen

Input & Output

Visualization

Time & Space Complexity

Related Problems

Constraints

Common Approaches

Brute Force (Check All Paths) — Algorithm Steps

Visualization

Code -

Time & Space Complexity

Constraints

Select Compiler