Swift Coding Test Course, Finding the Diameter of a Tree

Solving various algorithm problems for coding test preparation is very helpful in improving your programming skills. This time, we will take a deep dive into the ‘Finding the Diameter of a Tree’ problem. Understanding and knowing how to solve tree structures is very useful, as they frequently appear in many programming problems.

Problem Description

A tree is a non-linear data structure composed of nodes and edges. The diameter of a tree refers to the longest path between any two nodes in the tree. In other words, it is about finding the longest path between two nodes (the number of edges between the nodes). This problem can usually be solved using DFS (Depth-First Search) or BFS (Breadth-First Search).

Problem Definition

Given a tree, find the diameter of the tree.
Input: Number of nodes N in the tree (1 <= N <= 100,000)
      Edges in the form of N-1 pairs (u, v) representing each end node of the edge.
Output: Diameter of the tree
    

Problem Analysis

Due to the nature of trees, where all nodes are connected by exactly one path, the method for finding the longest distance between two nodes is to explore using either DFS or BFS. The general idea is as follows:

1. Choose an arbitrary node A and find the node B that is farthest from A.
2. Then, find the node C that is farthest from B. The distance at this point is the diameter of the tree.

Algorithm Implementation

Now let’s implement the above algorithm in Swift. We will use an adjacency list to represent the tree.

Swift Code Implementation

import Foundation

class Graph {
    var adjList: [[Int]]

    init(size: Int) {
        adjList = Array(repeating: [], count: size)
    }

    func addEdge(u: Int, v: Int) {
        adjList[u].append(v)
        adjList[v].append(u)
    }

    func bfs(start: Int) -> (Int, Int) {
        var dist = Array(repeating: -1, count: adjList.count)
        var queue: [Int] = []
        
        dist[start] = 0
        queue.append(start)

        var farthestNode = start
        var maxDistance = 0

        while !queue.isEmpty {
            let node = queue.removeFirst()
            for neighbor in adjList[node] {
                if dist[neighbor] == -1 { // If not visited
                    dist[neighbor] = dist[node] + 1
                    queue.append(neighbor)
                    if dist[neighbor] > maxDistance {
                        maxDistance = dist[neighbor]
                        farthestNode = neighbor
                    }
                }
            }
        }
        return (farthestNode, maxDistance)
    }
}

func findTreeDiameter(edges: [(Int, Int)], n: Int) -> Int {
    let graph = Graph(size: n)
    for edge in edges {
        graph.addEdge(u: edge.0 - 1, v: edge.1 - 1) // 0-indexed
    }

    let (farthestNode, _) = graph.bfs(start: 0)
    let (_, diameter) = graph.bfs(start: farthestNode)
    return diameter
}

// Example Input
let n = 5
let edges: [(Int, Int)] = [(1, 2), (1, 3), (2, 4), (2, 5)]
let diameter = findTreeDiameter(edges: edges, n: n)
print("Diameter of the tree: \(diameter)")
    

Code Explanation

This code calculates the diameter of the tree in the following way:

1. First, we define a Graph class to store the edge list.
2. We add edge information using the addEdge method.
3. The bfs function performs BFS starting from a given node and returns the farthest node and its distance.
4. The findTreeDiameter function creates a graph using the given edges and finds the diameter of the tree through two BFS calls.

Execution Result

When the code is executed, it produces the following output:

Diameter of the tree: 3

This indicates the longest path between node 4 and node 5.

Conclusion

In this lesson, we dealt with the problem of finding the diameter of a tree. This problem can be solved using DFS or BFS and helps enhance understanding of tree structures. It is a question that often appears in actual coding tests, so be sure to practice enough to become familiar with algorithm implementation. It is also beneficial to experiment with various tree cases for a deeper understanding.

In the next lesson, we will cover more algorithm problems, so please stay tuned!