Program to Connect a Forest in Python

Given a forest (collection of disconnected trees) represented as an adjacency list, we need to connect all trees into a single tree by adding the minimum number of edges. The goal is to minimize the diameter (longest path between any two nodes) of the resulting connected tree.

0 1 2 3 4 5 6 7 Add edge Forest ? Connected Tree

Algorithm Overview

The solution works in two phases:

  1. Find diameter of each tree: Calculate the longest path in each disconnected component
  2. Optimally connect trees: Connect trees to minimize the overall diameter

Finding Tree Diameter

For each tree, we use DFS to find the diameter (longest path between any two nodes) ?

import heapq
import math

def find_diameter(graph, start_node, seen):
    """Find diameter of a tree component starting from start_node"""
    diameter = 0
    
    def dfs(node, parent):
        nonlocal diameter
        seen.add(node)
        
        # Store the two longest paths from current node
        child_depths = []
        
        for neighbor in graph[node]:
            if neighbor != parent:
                depth = dfs(neighbor, node) + 1
                heapq.heappush(child_depths, depth)
                if len(child_depths) > 2:
                    heapq.heappop(child_depths)  # Keep only 2 largest
        
        if not child_depths:
            return 0
            
        # Update diameter with path through current node
        diameter = max(diameter, sum(child_depths))
        return max(child_depths)
    
    dfs(start_node, None)
    return diameter

# Test with a simple tree
test_graph = [
    [1, 2],    # Node 0 connects to 1, 2
    [0, 3, 4], # Node 1 connects to 0, 3, 4  
    [0],       # Node 2 connects to 0
    [1],       # Node 3 connects to 1
    [1]        # Node 4 connects to 1
]

seen = set()
diameter = find_diameter(test_graph, 0, seen)
print(f"Tree diameter: {diameter}")
Tree diameter: 3

Complete Solution

Here's the complete algorithm that connects all trees optimally ?

import heapq
import math

class ForestConnector:
    def connect_forest(self, graph):
        seen = set()
        tree_diameters = []
        max_single_diameter = 0
        
        # Find diameter of each disconnected tree
        for node in range(len(graph)):
            if node in seen:
                continue
                
            diameter = self.find_tree_diameter(graph, node, seen)
            max_single_diameter = max(max_single_diameter, diameter)
            tree_diameters.append(math.ceil(diameter / 2))
        
        # If only one tree, return its diameter
        if len(tree_diameters) <= 1:
            return max_single_diameter
        
        # Connect trees optimally
        return self.connect_optimally(tree_diameters, max_single_diameter)
    
    def find_tree_diameter(self, graph, start_node, seen):
        diameter = 0
        
        def dfs(node, parent):
            nonlocal diameter
            seen.add(node)
            
            child_depths = []
            for neighbor in graph[node]:
                if neighbor != parent:
                    depth = dfs(neighbor, node) + 1
                    heapq.heappush(child_depths, depth)
                    if len(child_depths) > 2:
                        heapq.heappop(child_depths)
            
            if not child_depths:
                return 0
                
            diameter = max(diameter, sum(child_depths))
            return max(child_depths)
        
        dfs(start_node, None)
        return diameter
    
    def connect_optimally(self, radii, max_single_diameter):
        # Find the largest radius and reduce it by 1
        max_radius = max(radii)
        for i in range(len(radii)):
            if radii[i] == max_radius:
                radii[i] -= 1
                break
        
        # Add 1 to each radius (cost of connecting)
        radii = [r + 1 for r in radii]
        
        # Find two largest radii after connection
        two_largest = heapq.nlargest(2, radii)
        connected_diameter = sum(two_largest)
        
        return max(connected_diameter, max_single_diameter)

# Test with the forest example
connector = ForestConnector()
graph = [
    [1, 2],    # Tree 1: 0-1-3, 0-1-4, 0-2
    [0, 3, 4], #         diameter = 3
    [0],
    [1],
    [1],
    [6, 7],    # Tree 2: 6-5-7
    [5],       #         diameter = 2  
    [5]
]

result = connector.connect_forest(graph)
print(f"Minimum possible diameter after connecting: {result}")
Minimum possible diameter after connecting: 4

How It Works

The algorithm follows these key steps:

  1. Calculate tree diameters: For each disconnected component, find its diameter using DFS
  2. Compute radii: The radius of a tree is ?diameter/2? (distance from center to farthest node)
  3. Optimal connection strategy: Connect trees through their centers to minimize the resulting diameter
  4. Final diameter: After connecting, the diameter is the maximum of the original largest diameter or the sum of the two largest radii plus connection costs

Time and Space Complexity

Aspect Complexity Explanation
Time O(V + E) DFS traversal of all nodes and edges
Space O(V) Recursion stack and seen set

Conclusion

This algorithm efficiently connects a forest into a single tree while minimizing the diameter. The key insight is connecting trees through their centers and choosing the optimal connection strategy based on tree radii.

Updated on: 2026-03-25T13:39:12+05:30

266 Views

Kickstart Your Career

Get certified by completing the course

Get Started
Advertisements