Skip to main content
Talentd

Pythagorean Distance Nodes in a Tree - Solution & Explanation

MediumTreeBreadth-First Search14 min readAsked at: PayPal
Practice this problem

Problem Statement

You are given an integer n and an undirected tree with n nodes numbered from 0 to n - 1. The tree is represented by a 2D array edges of length n - 1, where edges[i] = [ui, vi] indicates an undirected edge between ui and vi.

You are also given three distinct target nodes x, y, and z.

For any node u in the tree:

  • Let dx be the distance from u to node x
  • Let dy be the distance from u to node y
  • Let dz be the distance from u to node z

The node u is called special if the three distances form a Pythagorean Triplet.

Return an integer denoting the number of special nodes in the tree.

A Pythagorean triplet consists of three integers a, b, and c which, when sorted in ascending order, satisfy a2 + b2 = c2.

The distance between two nodes in a tree is the number of edges on the unique path between them.

 

Example 1:

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

Output: 3

Explanation:

For each node, we compute its distances to nodes x = 1, y = 2, and z = 3.

  • Node 0 has distances 1, 1, and 1. After sorting, the distances are 1, 1, and 1, which do not satisfy the Pythagorean condition.
  • Node 1 has distances 0, 2, and 2. After sorting, the distances are 0, 2, and 2. Since 02 + 22 = 22, node 1 is special.
  • Node 2 has distances 2, 0, and 2. After sorting, the distances are 0, 2, and 2. Since 02 + 22 = 22, node 2 is special.
  • Node 3 has distances 2, 2, and 0. After sorting, the distances are 0, 2, and 2. This also satisfies the Pythagorean condition.

Therefore, nodes 1, 2, and 3 are special, and the answer is 3.

Example 2:

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

Output: 0

Explanation:

For each node, we compute its distances to nodes x = 0, y = 3, and z = 2.

  • Node 0 has distances 0, 3, and 2. After sorting, the distances are 0, 2, and 3, which do not satisfy the Pythagorean condition.
  • Node 1 has distances 1, 2, and 1. After sorting, the distances are 1, 1, and 2, which do not satisfy the Pythagorean condition.
  • Node 2 has distances 2, 1, and 0. After sorting, the distances are 0, 1, and 2, which do not satisfy the Pythagorean condition.
  • Node 3 has distances 3, 0, and 1. After sorting, the distances are 0, 1, and 3, which do not satisfy the Pythagorean condition.

No node satisfies the Pythagorean condition. Therefore, the answer is 0.

Example 3:

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

Output: 1

Explanation:

For each node, we compute its distances to nodes x = 1, y = 3, and z = 0.

  • Node 0 has distances 1, 2, and 0. After sorting, the distances are 0, 1, and 2, which do not satisfy the Pythagorean condition.
  • Node 1 has distances 0, 1, and 1. After sorting, the distances are 0, 1, and 1. Since 02 + 12 = 12, node 1 is special.
  • Node 2 has distances 1, 2, and 2. After sorting, the distances are 1, 2, and 2, which do not satisfy the Pythagorean condition.
  • Node 3 has distances 1, 0, and 2. After sorting, the distances are 0, 1, and 2, which do not satisfy the Pythagorean condition.

Therefore, the answer is 1.

 

Constraints:

  • 4 <= n <= 105
  • edges.length == n - 1
  • edges[i] = [ui, vi]
  • 0 <= ui, vi, x, y, z <= n - 1
  • x, y, and z are pairwise distinct.
  • The input is generated such that edges represent a valid tree.

Approach Overview

Problem Overview: You are given a tree and need to identify nodes whose distances participate in a Pythagorean relationship. If three node distances a, b, and c satisfy a² + b² = c², they form a valid Pythagorean triple. The task reduces to computing distances in the tree and checking combinations that satisfy this equation.

Approach 1: All-Pairs Distance Enumeration (O(n²) time, O(n) space)

A direct strategy computes distances between many node pairs and checks whether any triple satisfies the Pythagorean condition. Because the graph is a tree, you can run a BFS from each node to compute distances to all others. For every pair of nodes with distances a and b, compute a² + b² and check whether another node exists with distance sqrt(a²+b²). This approach is straightforward but expensive since running BFS repeatedly leads to roughly O(n²) total work.

Approach 2: BFS + Distance Enumeration (O(n²) time, O(n) space)

The optimized method runs a single BFS from a chosen root to compute the distance (depth) of every node. Store these distances in an array and optionally keep them in a hash set for constant-time lookup. Then enumerate all pairs of depths (a, b). For each pair, compute c² = a² + b² and check whether c is an integer and exists among the recorded distances. Because BFS computes distances in O(n) time and pair enumeration costs O(n²), the overall complexity remains manageable for typical constraints.

The key insight is that a tree has exactly one path between nodes. Running BFS once gives correct shortest distances from the root, turning the problem into a numeric check over distances rather than repeated graph traversals. Using a hash set for depth lookup removes the need for extra scanning.

Problems like this commonly combine traversal with arithmetic constraints. Mastering Breadth-First Search and understanding how distances behave in a tree lets you reduce the graph portion to a preprocessing step, after which the remaining work is simple enumeration and math checks.

Recommended for interviews: The BFS + enumeration approach. Interviewers expect you to recognize that computing distances in a tree is easiest with BFS, then reduce the rest of the problem to checking the Pythagorean equation. Mentioning the brute-force multi-BFS idea shows baseline understanding, but the single-BFS solution demonstrates stronger algorithmic thinking and familiarity with graph traversal patterns.

Solution

We first construct an adjacency list g based on the edges given in the problem, where g[u] stores all nodes adjacent to node u.

Next, we define a function bfs(i) to calculate the distances from node i to all other nodes. We use a queue to implement Breadth-First Search (BFS) and maintain a distance array dist, where dist[j] represents the distance from node i to node j. Initially, dist[i] = 0, and the distances to all other nodes are set to infinity. During the BFS process, we continuously update the distance array until all reachable nodes have been traversed.

We call bfs(x), bfs(y), and bfs(z) to calculate the distances from nodes x, y, and z to all other nodes, obtaining three distance arrays d_1, d_2, and d_3 respectively.

Finally, we iterate through all nodes u. For each node, we retrieve its distances to x, y, and z as a = d_1[u], b = d_2[u], and c = d_3[u]. We sort these three distances and check if they satisfy the Pythagorean theorem condition: a^2 + b^2 = c^2. If the condition is met, we increment the answer count.

The time complexity is O(n), and the space complexity is O(n), where n is the number of nodes in the tree.

Code

PythonJavaC++GoTypeScript

Python

Java

C++

Go

TypeScript

Try this approach in the editor →

Detailed Complexity Analysis

ApproachTimeSpaceWhen to Use
All-Pairs BFS Distance EnumerationO(n²)O(n)Conceptual baseline when exploring tree distance relationships
Single BFS + Pair EnumerationO(n²)O(n)Best practical solution; compute depths once and check Pythagorean pairs efficiently

Video Solution

Pythagorean Distance Nodes in a Tree 🔥 LeetCode 3820 | Weekly Contest 486 | Tree + Math • Study Placement • 501 views views

Watch 9 more video solutions →

Frequently Asked Questions

Is Pythagorean Distance Nodes in a Tree easy or hard?
The problem is generally considered Medium difficulty. BFS on a tree is straightforward, but recognizing how to convert the graph distances into a Pythagorean triple check and implementing efficient enumeration requires careful reasoning.
Pythagorean Distance Nodes in a Tree Python/Java solution
Typical implementations compute node depths using BFS, store them in an array or list, and use a set for quick membership checks. The rest of the logic iterates through depth pairs and verifies the Pythagorean equation using integer square root checks. This pattern translates directly to Python, Java, C++, Go, and TypeScript.
How to solve Pythagorean Distance Nodes in a Tree in O(n)?
Pure O(n) solutions are generally not feasible because verifying the Pythagorean condition requires comparing multiple distance combinations. The practical approach is O(n²): run BFS to compute depths, store them in a set, and check every pair (a, b) to see if sqrt(a² + b²) exists among the depths.
What is the best approach for Pythagorean Distance Nodes in a Tree?
The most practical approach is running a single BFS to compute distances (depths) from a root node, then enumerating pairs of distances and checking whether they satisfy a² + b² = c². A hash set of depths allows constant-time lookups when checking whether the required third distance exists. This keeps the traversal cost at O(n) and the enumeration cost at O(n²).
Is Pythagorean Distance Nodes in a Tree asked at Google/Amazon/Meta?
Problems combining tree traversal with numeric constraints appear frequently in interviews at companies like Google, Amazon, and Meta. Variants often test BFS or DFS for distance computation along with mathematical checks or hash-based lookups.
What data structure is used in Pythagorean Distance Nodes in a Tree?
The solution primarily uses an adjacency list to represent the tree and a queue for BFS traversal. A hash set or hash map stores node depths so that potential Pythagorean matches can be checked in constant time.
What is the time complexity of Pythagorean Distance Nodes in a Tree?
The common solution runs BFS once to compute distances in O(n) time and then enumerates all pairs of node distances in O(n²). Checking whether a² + b² forms a valid square is constant time with a hash lookup. Overall complexity is O(n²) time and O(n) space.

Ready to solve this problem?

Practice Pythagorean Distance Nodes in a Tree with our built-in code editor and test cases.

Practice on FleetCode
Previous

Binary Tree Inorder Traversal

Next

Unique Binary Search Trees II