Python Coding Test Course, Finding the Diameter of a Tree

Trees are one of the most important data structures in computer science. In particular, trees are useful for representing hierarchical relationships and are used in various algorithmic problems. This lecture will cover the problem of finding the diameter of a tree.
The diameter refers to the longest path between two nodes in the tree. This problem can be solved using DFS (Depth First Search) or BFS (Breadth First Search) algorithms.

Problem Description

Each node in the given non-empty tree is represented by an integer. Solve the problem of finding the length of the longest path between two nodes in the tree.
The input consists of the number of nodes in the tree N and N-1 edge information. The edge information is provided in a way that connects two nodes.
Specifically, the input will be given in the following format: 

    N
    a1 b1
    a2 b2
    ...
    a(N-1) b(N-1)

Here, a and b represent the two connected nodes, respectively.

Input Example

    5
    1 2
    1 3
    2 4
    2 5

Output Example

    3

In this case, the diameter of the tree is between node 4 and node 5, with the path being 4 → 2 → 1 → 3 or 4 → 2 → 5.
Therefore, the output is 3.

Solution

To find the diameter of the tree, we can use DFS or BFS algorithms.
The general approach is as follows:

  1. In the first step, perform DFS from an arbitrary node to find the farthest node.
    Let’s call this node X.
  2. In the second step, perform DFS again from node X to find the farthest node, Y.
    The path between X and Y will be the diameter of the tree.

Through this process, the time complexity will be O(N), implemented by recursively calling DFS.

Python Code Implementation

Now, based on the logic above, let’s implement the code to find the diameter of the tree in Python.
Check the details of each step with the code provided below. 

from collections import defaultdict

def dfs(graph, node, visited):
    visited.add(node)
    max_distance = 0
    farthest_node = node

    for neighbor in graph[node]:
        if neighbor not in visited:
            distance, next_node = dfs(graph, neighbor, visited)
            distance += 1
            
            if distance > max_distance:
                max_distance = distance
                farthest_node = next_node

    return max_distance, farthest_node

def tree_diameter(edges, n):
    graph = defaultdict(list)
    
    for a, b in edges:
        graph[a].append(b)
        graph[b].append(a)

    # Step 1: start DFS from an arbitrary node (1)
    visited = set()
    _, farthest_node = dfs(graph, 1, visited)

    # Step 2: start DFS from the farthest node found
    visited.clear()
    diameter, _ = dfs(graph, farthest_node, visited)

    return diameter

# Input reading part
n = int(input())
edges = [tuple(map(int, input().split())) for _ in range(n-1)]
print(tree_diameter(edges, n))

Code Explanation

The above code is structured in the following way:

  • collections.defaultdict is used to create the graph in the form of an adjacency list.
    This represents the connectivity between nodes.
  • dfs function performs depth-first search and calculates the distance to each node.
    It returns the farthest node and distance.
  • tree_diameter function coordinates the overall process and calculates the diameter through two DFS calls.
  • In the last part, it takes input from the user and calls the tree_diameter function to output the result.

Performance Analysis

The presented algorithm has a time complexity of O(N).
This is possible because it visits all nodes in the tree once through DFS.
Therefore, it can efficiently calculate the diameter even for very large trees.

Conclusion

In this lecture, we explored the diameter of trees.
We were able to efficiently solve the problem using a DFS approach.
Trees are utilized in various problems, so it is beneficial to thoroughly understand the contents of this lecture.
If you have additional questions or need practice problems, please leave a comment.