Tree Diameter II

# Tree Diameter II ## Problem Statement You are given a tree consisting of $n$ nodes, numbered $1, 2, \dots, n$. The distance between two nodes in a tree is defined as the number of edges on the unique simple path between them. The diameter of a tree is the maximum distance between any pair of nodes in the tree. Your task is to count the number of distinct pairs of nodes $(u, v)$ such that the distance between $u$ and $v$ is equal to the tree's diameter. Note that pairs $(u, v)$ and $(v, u)$ are considered the same diameter and should be counted once. ## Input Format - The first line contains a single integer $n$ ($1 \le n \le 2 \times 10^5$), representing the number of nodes in the tree. - The next $n-1$ lines describe the edges. Each line contains two integers $u$ and $v$ ($1 \le u, v \le n$), indicating an edge between nodes $u$ and $v$. ## Output Format Print a single integer: the count of distinct pairs of nodes that form a diameter of the tree. ## Sample Test Cases ### Sample Input 1 ``` 5 1 2 1 3 3 4 3 5 ``` ### Sample Output 1 ``` 2 ``` ### Sample Input 2 ``` 4 1 2 2 3 3 4 ``` ### Sample Output 2 ``` 1 ``` ## Constraints - $1 \le n \le 2 \times 10^5$ - $1 \le u, v \le n$ - The given graph is guaranteed to be a tree. - Time limit: 1 second - Memory limit: 256 MB ## Explanation ### Sample 1 The tree has $n=5$ nodes and edges $(1,2)$, $(1,3)$, $(3,4)$, $(3,5)$. The diameter of this tree is $3$. The pairs of nodes that achieve this maximum distance are $(2,4)$ and $(2,5)$. - The path $2 \to 1 \to 3 \to 4$ has length $3$. - The path $2 \to 1 \to 3 \to 5$ has length $3$. Thus, there are $2$ distinct pairs of nodes that form a diameter. ### Sample 2 The tree has $n=4$ nodes, forming a path $1 \to 2 \to 3 \to 4$. The diameter of this tree is $3$. The only pair of nodes that achieves this maximum distance is $(1,4)$. - The path $1 \to 2 \to 3 \to 4$ has length $3$. Thus, there is $1$ distinct pair of nodes that forms a diameter. Solution Approach The solution performs a DFS that for each node returns a pair: the maximum depth from that node to a leaf (measured in number of nodes) and the number of distinct nodes at that maximum depth in its subtree. For a node, collect the pairs from all its children and sort them by depth descending. If there are no children, the node contributes a path of length 1 1 with count 1 1. If there is exactly one child, potential diameters passing through this node are just the child depth plus one, with the same count as that child. If there are at least two children, consider the top two depths. If the top two depths are equal, any pair of nodes coming from different children with that depth will form a diameter passing through the current node; count combinations accordingly (sum over children counts at that depth, then number of unordered pairs). If the top two depths differ, take the product of counts from children with the top and second-top depths. Each time a candidate diameter length (measured in number of nodes on the path) is found, update the global best diameter and add up counts if the same diameter occurs elsewhere. The algorithm is implemented exactly as in the editorial code: a DFS that returns ( maxDepthPlusOne, countAtMax ) (maxDepthPlusOne, countAtMax) and a global update routine that manages diameter diameter and cnt_diameter cnt_diameter. Time Complexity per test case: O ( n log ⁡ n ) O(nlogn) in the worst case (sorting children lists at each node can lead to overall ∑ deg ⁡ ( v ) log ⁡ deg ⁡ ( v ) ∑deg(v)logdeg(v) which is O ( n log ⁡ n ) O(nlogn) in the worst case). Space Complexity per test case: O ( n ) O(n) for the adjacency list and recursion stack.