Count Tree Vertices with Pythagorean Distances

Given an unweighted, undirected tree with n vertices labeled 0 through n - 1, and three distinct fixed vertices x, y, and z, count the vertices whose distances to those three fixed vertices form a positive Pythagorean triple.

For each vertex v, sort its three distances into a <= b <= c. Count v exactly once if a > 0 and a * a + b * b == c * c. A distance is the number of edges on the unique path between two vertices.

Input

• n: the number of vertices.
• edges: n - 1 pairs of vertex labels describing a connected tree.
• x, y, z: three distinct vertex labels in the tree.

Output

Return the number of qualifying vertices as an integer.

Constraints and Edge Cases

• For this practice version, 3 <= n <= 100000.
• Treat edges as unweighted and the three fixed vertices as distinct in this practice version.
• Every edge joins two different vertices. There are no duplicate edges.
• The tree may be a chain; do not assume its depth is small.
• The three distances have no required association with x, y, and z after sorting.
• A vertex at distance zero from a fixed vertex cannot qualify, even if its squared distances otherwise satisfy the equation.
• Square distances using arithmetic that can represent values up to (n - 1) * (n - 1) exactly.

Example 1

n = 12
edges = [[0,1],[1,2],[2,3],[0,4],[4,5],[5,6],[6,7],[0,8],[8,9],[9,10],[10,11]]
x = 3
y = 7
z = 11
output = 2

Vertices 4 and 8 each have sorted distances [3, 4, 5]. No other vertex qualifies.

Example 2

n = 3
edges = [[0,1],[1,2]]
x = 0
y = 1
z = 2
output = 0

Every vertex is one of the fixed vertices and therefore has a zero distance.

Count Tree Vertices with Pythagorean Distances

Approach

Build an undirected adjacency list, then run an iterative breadth-first traversal from each fixed vertex. Each traversal stores the distance of every vertex and marks a vertex when it first enters the queue. Use an array or list with a moving head index, so removing the next queue item does not shift the remaining items.

For each vertex, order its three recorded distances. Reject it when the smallest is zero; otherwise compare the sum of the two smaller squares with the largest square and increment the answer only on exact equality. Sorting three numbers is constant work; computing their minimum, maximum and middle value is equivalent.

A breadth-first traversal visits vertices in nondecreasing distance from its source. When an unseen neighbor is reached, extending the current path adds one edge and gives its shortest distance. The tree is connected, so each traversal reaches every vertex. Thus the three arrays hold precisely the distances in the statement. The final test applies exactly the positivity and Pythagorean conditions to those three values, independent of which fixed vertex supplies the longest distance. Every vertex is considered once, so none is omitted or counted twice. No primitive-triple restriction is imposed.

There are n vertices and n-1 edges. Three traversals and one final scan take O(n) time; adjacency lists, distances and a queue take O(n) space. All traversal is iterative, including for a chain of 100000 vertices. Distances themselves fit in 32 bits, but Java and C++ convert them to 64-bit variables before multiplication or addition: a square may be 9999800001 and the sum may be 19999600002. Python integers and JavaScript Number represent these integers exactly; no 32-bit bitwise coercion is used. The input edges are not modified.

• Time complexity: O(n): three traversals of a tree and one constant-work scan per vertex.
• Space complexity: O(n) for adjacency lists, distance arrays and iterative queues.