Program to find minimum cost to connect all points in Python

Suppose we have an array called points with some coordinates in the form (x, y). The cost of connecting two points (xi, yi) and (xj, yj) is the Manhattan distance between them, calculated as |xi - xj| + |yi - yj|. We need to find the minimum cost to connect all points using a Minimum Spanning Tree (MST) approach.

So, if the input is like points = [(0,0),(3,3),(2,10),(6,3),(8,0)], then the output will be 22.

(0,0) (3,3) (2,10) (6,3) (8,0) 6 5 3 8

The total minimum distance is (6+5+3+8) = 22.

Algorithm: Prim's MST

We use Prim's algorithm with a min-heap to find the MST ?

  • points_set := a set holding indices from 0 to size of points - 1
  • heap := initialize with (0, 0) representing (distance, index)
  • visited_node := track visited points
  • total_distance := accumulate minimum cost
  • While heap is not empty and not all points visited:
    • Pop minimum distance edge from heap
    • If point not visited, add it to MST and explore neighbors
    • Add Manhattan distances to unvisited points to heap

Example

Let us see the implementation using Prim's algorithm ?

import heapq

def solve(points):
    points_set = set(range(len(points)))
    heap = [(0, 0)]  # (distance, index)
    visited_node = set()
    total_distance = 0
    
    while heap and len(visited_node) < len(points):
        distance, current_index = heapq.heappop(heap)
        
        if current_index not in visited_node:
            visited_node.add(current_index)
            points_set.discard(current_index)
            total_distance += distance
            
            x0, y0 = points[current_index]
            
            # Add edges to all unvisited points
            for next_index in points_set:
                x1, y1 = points[next_index]
                manhattan_dist = abs(x0 - x1) + abs(y0 - y1)
                heapq.heappush(heap, (manhattan_dist, next_index))
    
    return total_distance

# Test the function
points = [(0,0),(3,3),(2,10),(6,3),(8,0)]
result = solve(points)
print("Minimum cost to connect all points:", result)
Minimum cost to connect all points: 22

How It Works

The algorithm works by:

  • Starting from point 0 with cost 0
  • Expanding the MST by always choosing the minimum cost edge
  • Using a heap to efficiently get the next minimum edge
  • Avoiding cycles by checking if points are already visited

Time Complexity

The time complexity is O(N² log N) where N is the number of points, since we potentially add O(N²) edges to the heap and each heap operation takes O(log N) time.

Conclusion

This problem is solved using Prim's MST algorithm with Manhattan distance. The min-heap ensures we always select the minimum cost edge, resulting in the optimal solution for connecting all points.

Updated on: 2026-03-26T13:47:58+05:30

530 Views

Kickstart Your Career

Get certified by completing the course

Get Started
Advertisements