We use cookies (including Google cookies) to personalize ads and analyze traffic. By continuing to use our site, you accept our Privacy Policy.

SimplyLeet
TagsDisclaimersPrivacyTermsContactAbout Us

Count Valid Paths in a Tree

Difficulty: Hard

Problem Description

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.

Key Insights

  • A tree is a connected acyclic graph, which means there is a unique path between any two nodes.
  • We need to count paths that contain exactly one prime labeled node.
  • Efficiently determining the primality of node labels is crucial since n can be as large as 100,000.
  • A depth-first search (DFS) can be employed to explore paths and count valid ones by tracking the number of prime nodes encountered.

Space and Time Complexity

Time Complexity: O(n) - each node and edge is processed once. Space Complexity: O(n) - for storing the adjacency list representation of the tree.

Solution

We will use a Depth-First Search (DFS) approach to traverse the tree. The algorithm will maintain a count of prime nodes encountered along the path from the root to each node. Whenever we reach a node, we will check the paths that can be formed with the current node as the endpoint and count those that have exactly one prime node. The tree will be represented using an adjacency list for efficient traversal.

  1. Build the tree using an adjacency list from the edges.
  2. Determine prime numbers up to n using the Sieve of Eratosthenes.
  3. Use DFS to traverse the tree, maintaining a count of prime nodes in the current path.
  4. For each node, check possible paths that end at that node and count how many have exactly one prime.

Code Solutions

def count_valid_paths(n, edges):
    from collections import defaultdict

    # Function to determine prime numbers using Sieve of Eratosthenes
    def sieve(n):
        is_prime = [True] * (n + 1)
        is_prime[0] = is_prime[1] = False
        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
        return is_prime

    # Create adjacency list for the tree
    tree = defaultdict(list)
    for u, v in edges:
        tree[u].append(v)
        tree[v].append(u)

    # Get prime numbers up to n
    is_prime = sieve(n)

    # Variable to store the count of valid paths
    valid_paths = 0

    # DFS function to traverse the tree
    def dfs(node, parent, prime_count):
        nonlocal valid_paths
        # If current node is prime, increment the count
        if is_prime[node]:
            prime_count += 1

        # Check valid paths with current node
        if prime_count == 1:
            valid_paths += 1

        # Continue DFS for child nodes
        for neighbor in tree[node]:
            if neighbor != parent:  # Avoid going back to parent
                dfs(neighbor, node, prime_count)

    # Start DFS from node 1 (or any arbitrary node) with no parent and initial prime count
    dfs(1, -1, 0)

    return valid_paths

# Example usage
print(count_valid_paths(5, [[1,2],[1,3],[2,4],[2,5]]))  # Output: 4
← Back to All Questions