DEV Community

Cover image for LeetCode, Hard (Acceptance 24%, Latest): 2867. Count Valid Paths in a Tree. DFS. O(n). Swift.
Sergey Leschev
Sergey Leschev

Posted on

LeetCode, Hard (Acceptance 24%, Latest): 2867. Count Valid Paths in a Tree. DFS. O(n). Swift.

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] = [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.
Enter fullscreen mode Exit fullscreen mode

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.
Enter fullscreen mode Exit fullscreen mode

Constraints:
1 <= n <= 10^5
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.

Intuition

The intuition is to employ a depth-first search (DFS) approach through the tree.

Approach

During each step in the traversal, we perform the following key calculations:

  1. Determine the path that ends at the current node.
  2. Compute two different subtree paths that traverse the current node.
  3. Maintain an array that keeps track of the cases where paths contain either 0 or 1 prime number.

This method allows us to efficiently count the valid paths in the tree while considering the presence of prime numbers.

Complexity

  • Time complexity: O(n)
  • Space complexity: O(n)

Code (Swift)

class Solution {
    func countPaths(_ n: Int, _ edges: [[Int]]) -> Int {
        var ans = 0
        var primes = Array(repeating: 1, count: n + 5)

        primes[1] = 0
        for i in 2...Int(sqrt(Double(n + 5))) {
            if primes[i] == 0 {
                continue
            }
            for j in stride(from: i * i, to: primes.count, by: i) {
                primes[j] = 0
            }
        }

        var adjList = Array(repeating: [Int](), count: n + 1)
        for edge in edges {
            adjList[edge[0]].append(edge[1])
            adjList[edge[1]].append(edge[0])
        }

        func dfs(_ curr: Int, _ parent: Int, _ adjList: [[Int]]) -> [Int] {
            var count = [0, 0]
            let isPrime = primes[curr] == 1

            for child in adjList[curr] {
                if child == parent {
                    continue
                }
                let next = dfs(child, curr, adjList)
                if isPrime {
                    // Path ends at curr
                    ans += next[0]
                    // curr is conduit for path
                    ans += count[1] * next[0]

                    count[1] += next[0]
                } else {
                    // Ends at curr
                    ans += next[1]
                    // curr is conduit for path
                    ans += count[0] * next[1]
                    ans += count[1] * next[0]

                    count[0] += next[0]
                    count[1] += next[1]
                }
            }

            count[isPrime ? 1 : 0] += 1
            return count
        }

        dfs(1, -1, adjList)
        return ans
    }
}
Enter fullscreen mode Exit fullscreen mode

Sources: Github


Contacts
I have a clear focus on time-to-market and don't prioritize technical debt. And I took part in the Pre-Sale/RFX activity as a System Architect, assessment efforts for Mobile (iOS-Swift, Android-Kotlin), Frontend (React-TypeScript) and Backend (NodeJS-.NET-PHP-Kafka-SQL-NoSQL). And I also formed the work of Pre-Sale as a CTO from Opportunity to Proposal via knowledge transfer to Successful Delivery.

🛩ī¸ #startups #management #cto #swift #typescript #database
📧 Email: sergey.leschev@gmail.com
👋 LinkedIn: https://linkedin.com/in/sergeyleschev/
👋 LeetCode: https://leetcode.com/sergeyleschev/
👋 Twitter: https://twitter.com/sergeyleschev
👋 Github: https://github.com/sergeyleschev
🌎 Website: https://sergeyleschev.github.io
🌎 Reddit: https://reddit.com/user/sergeyleschev
🌎 Quora: https://quora.com/sergey-leschev
🌎 Medium: https://medium.com/@sergeyleschev
🖨ī¸ PDF Design Patterns: Download

Top comments (1)

Collapse
 
overflow profile image
Info Comment hidden by post author - thread only accessible via permalink
overFlow

YOU !!!! You mean business!!!! When I grow up I wanna be like you.
I keep seeing your stuff!! although they are outside my scope of things I cannot waiti until the day i can understand them...
Hard ++ ??? if you were a sildenafil You would be dangerous....

Some comments have been hidden by the post's author - find out more