Count Valid Paths in a Tree

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.

Constraints

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

Visualization

Tap to expand

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

✓ Brute Force (Check All Paths)

⏱️ Time: O(n³) Space: O(n)

The naive approach generates every possible path between any two nodes in the tree and checks if exactly one prime number exists in each path. This involves finding all pairs of nodes and then traversing the path between them to count prime numbers.

DFS with Tree Properties (Optimized)

⏱️ Time: O(n²) Space: O(n)

This approach uses the fundamental property that in a tree, there's exactly one path between any two nodes. We perform DFS from each node and count valid paths more efficiently by avoiding redundant path calculations and using dynamic programming concepts.

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;

int findPath(int start, int end, AdjList* graph, bool* visited, int* path, int pathLen);

long long countPaths(int n, int** edges, int edgesSize) {
    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];
        if (graph[u].size >= graph[u].capacity) {
            graph[u].capacity *= 2;
            graph[u].adj = (int*)realloc(graph[u].adj, graph[u].capacity * sizeof(int));
        }
        if (graph[v].size >= graph[v].capacity) {
            graph[v].capacity *= 2;
            graph[v].adj = (int*)realloc(graph[v].adj, graph[v].capacity * sizeof(int));
        }
        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;
}

int main() {
    int n;
    scanf("%d", &n);
    
    char edgesStr[10000];
    scanf("%s", edgesStr);
    
    int** edges = NULL;
    int edgesSize = 0;
    
    if (strcmp(edgesStr, "[]") != 0) {
        // Simple parsing for format [[1,2],[1,3],[2,4]]
        edges = (int**)malloc(100 * sizeof(int*));
        for (int i = 0; i < 100; i++) {
            edges[i] = (int*)malloc(2 * sizeof(int));
        }
        
        int len = strlen(edgesStr);
        for (int i = 0; i < len; i++) {
            if (edgesStr[i] == '[' && i > 0) {
                int j = i + 1;
                while (j < len && edgesStr[j] != ']') j++;
                
                // Extract numbers
                int u = 0, v = 0;
                int k = i + 1;
                while (k < j && edgesStr[k] != ',') {
                    u = u * 10 + (edgesStr[k] - '0');
                    k++;
                }
                k++; // skip comma
                while (k < j) {
                    v = v * 10 + (edgesStr[k] - '0');
                    k++;
                }
                
                edges[edgesSize][0] = u;
                edges[edgesSize][1] = v;
                edgesSize++;
                i = j;
            }
        }
    }
    
    printf("%lld\n", countPaths(n, edges, edgesSize));
    
    if (edges) {
        for (int i = 0; i < 100; i++) {
            free(edges[i]);
        }
        free(edges);
    }
    
    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

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

Constraints

Visualization

Related Problems

Common Approaches

Brute Force (Check All Paths) — Algorithm Steps

Visualization

Code -

Time & Space Complexity

Select Compiler