LeetCode 3820 - Pythagorean Distance Nodes in a Tree

The problem requires identifying special nodes in a tree based on distances to three distinct target nodes, x, y, and z.

leetcodemediumtreebreadth-first-search

LeetCode Problem 3820

Difficulty: 🟡 Medium
Topics: Tree, Breadth-First Search

Solution

Problem Understanding

The problem requires identifying special nodes in a tree based on distances to three distinct target nodes, x, y, and z. A node is special if the distances from it to x, y, and z form a Pythagorean triplet, meaning that when sorted in ascending order, the sum of squares of the two smaller distances equals the square of the largest distance.

The input provides n nodes and a list of edges describing an undirected tree. Each edge connects two nodes, and because the input is guaranteed to be a tree, there is exactly one path between any two nodes. The goal is to compute distances from each node to x, y, and z, check if these distances form a Pythagorean triplet, and count the number of nodes that satisfy this condition.

Key points are that n can be as large as 10^5, so any solution that computes distances in O(n^2) per node will be too slow. The constraints guarantee the graph is a valid tree, nodes are numbered from 0 to n-1, and the target nodes x, y, z are distinct.

Edge cases to consider include nodes being one of the target nodes themselves (distance 0), all distances being equal, or distances forming degenerate Pythagorean triplets such as (0, a, a).

Approaches

Brute Force

The brute force method would compute the shortest path distance from every node u to each of the three target nodes x, y, and z. For each node, you would perform a BFS or DFS three times to calculate dx, dy, and dz, then check if (dx, dy, dz) forms a Pythagorean triplet. This approach works but has time complexity O(n^2), which is too slow for n up to 10^5.

Optimal Approach

The key insight is that distances in a tree can be efficiently computed using BFS from each target node. By performing three BFS traversals, one from x, one from y, and one from z, we can precompute the distances from all nodes to each target in O(n) time per BFS. Then, for each node, we only need to check the Pythagorean condition using the precomputed distances. This reduces the total time complexity to O(n), which is feasible for large n.

Approach Time Complexity Space Complexity Notes
Brute Force O(n^2) O(n) BFS from every node to x, y, z individually
Optimal O(n) O(n) BFS from each of x, y, z and then check distances

Algorithm Walkthrough

  1. Construct an adjacency list from the edges for efficient tree traversal.
  2. Initialize three arrays distX, distY, and distZ to store distances from each node to x, y, and z, respectively.
  3. Perform BFS starting from node x to compute the distance from x to every other node. Store the results in distX.
  4. Repeat BFS starting from y and z to compute distY and distZ.
  5. Iterate through each node in the tree. For a node u, retrieve the distances (dx, dy, dz) from the arrays.
  6. Sort the distances and check if they satisfy the Pythagorean triplet condition: a^2 + b^2 == c^2.
  7. Increment a counter for each node that satisfies the condition.
  8. Return the counter as the final number of special nodes.

Why it works: In a tree, BFS guarantees the shortest distance from a starting node to all other nodes. Precomputing distances in this way ensures that every node's distances to x, y, z are accurate. The sorting step ensures the correct application of the Pythagorean condition.

Python Solution

from collections import deque
from typing import List

class Solution:
    def specialNodes(self, n: int, edges: List[List[int]], x: int, y: int, z: int) -> int:
        def bfs(start: int) -> List[int]:
            dist = [-1] * n
            dist[start] = 0
            queue = deque([start])
            while queue:
                node = queue.popleft()
                for nei in adj[node]:
                    if dist[nei] == -1:
                        dist[nei] = dist[node] + 1
                        queue.append(nei)
            return dist
        
        adj = [[] for _ in range(n)]
        for u, v in edges:
            adj[u].append(v)
            adj[v].append(u)
        
        distX = bfs(x)
        distY = bfs(y)
        distZ = bfs(z)
        
        count = 0
        for i in range(n):
            a, b, c = sorted([distX[i], distY[i], distZ[i]])
            if a * a + b * b == c * c:
                count += 1
        
        return count

The Python code follows the algorithm closely. It uses BFS to calculate distances efficiently, stores them in separate arrays, and then evaluates the Pythagorean condition for every node. Sorting distances ensures correctness regardless of their relative order.

Go Solution

package main

func specialNodes(n int, edges [][]int, x int, y int, z int) int {
    adj := make([][]int, n)
    for _, e := range edges {
        u, v := e[0], e[1]
        adj[u] = append(adj[u], v)
        adj[v] = append(adj[v], u)
    }

    bfs := func(start int) []int {
        dist := make([]int, n)
        for i := range dist {
            dist[i] = -1
        }
        dist[start] = 0
        queue := []int{start}
        for len(queue) > 0 {
            node := queue[0]
            queue = queue[1:]
            for _, nei := range adj[node] {
                if dist[nei] == -1 {
                    dist[nei] = dist[node] + 1
                    queue = append(queue, nei)
                }
            }
        }
        return dist
    }

    distX := bfs(x)
    distY := bfs(y)
    distZ := bfs(z)

    count := 0
    for i := 0; i < n; i++ {
        a, b, c := distX[i], distY[i], distZ[i]
        if a > b { a, b = b, a }
        if b > c { b, c = c, b }
        if a > b { a, b = b, a }
        if a*a + b*b == c*c {
            count++
        }
    }

    return count
}

In Go, we implement BFS similarly and use slice-based queues. Sorting the distances is done manually using conditional swaps, avoiding the need for a separate sorting function. The logic otherwise mirrors the Python implementation.

Worked Examples

Example 1: n = 4, edges = [[0,1],[0,2],[0,3]], x=1, y=2, z=3

Node dx dy dz Sorted Pythagorean?
0 1 1 1 1,1,1 No
1 0 2 2 0,2,2 Yes
2 2 0 2 0,2,2 Yes
3 2 2 0 0,2,2 Yes

Answer: 3

Example 2: n = 4, edges = [[0,1],[1,2],[2,3]], x=0, y=3, z=2

All nodes fail the Pythagorean condition; answer is 0.

Example 3: n = 4, edges = [[0,1],[1,2],[1,3]], x=1, y=3, z=0

Only node 1 satisfies the Pythagorean condition; answer is 1.

Complexity Analysis

Measure Complexity Explanation
Time O(n) BFS is O(n) per target, three targets give O(3n) = O(n). Checking distances is O(n).
Space O(n) Distance arrays and adjacency list take O(n) space. Queue in BFS is O(n) in worst case.

The algorithm is efficient and scales to the maximum constraint n=10^5 because BFS and simple arithmetic checks are linear.

Test Cases

# Provided examples
assert Solution().specialNodes(4, [[0,1],[0,2],[0,3]], 1, 2, 3) == 3  # Example 1
assert Solution().specialNodes(4, [[0,1],[1,2],[2,3]], 0, 3, 2) == 0  # Example 2
assert Solution().specialNodes(4, [[0,1],[1,2],[1,3]], 1, 3, 0) == 1  # Example 3

# Edge cases
assert Solution().specialNodes(4, [[0,