min-cost-to-connect-all-points.py (1253B)
1 # Uses minimum spanning tree, induced using prim's algorith 2 # Uses an adjacency list to represent edges and costs 3 4 def man_dist(p_1, p_2): 5 return abs(p_1[0] - p_2[0]) + abs(p_1[1] - p_2[1]) 6 7 class Solution: 8 def minCostConnectPoints(self, points: List[List[int]]) -> int: 9 10 edges = {} 11 12 for p_1 in range(len(points) - 1): 13 for p_2 in range(p_1 + 1, len(points)): 14 if not p_1 in edges: 15 edges[p_1] = {} 16 if not p_2 in edges: 17 edges[p_2] = {} 18 19 dst = man_dist(points[p_1], points[p_2]) 20 edges[p_1][p_2] = dst 21 edges[p_2][p_1] = dst 22 23 24 # starting at point 0. 25 used = set() 26 used.add(0) 27 dist = 0 28 29 while len(used) != len(points): 30 31 best_dist = -1 32 best_edge = -1 33 34 for i in used: 35 for edge in edges[i]: 36 if not edge in used and (best_dist == -1 or edges[i][edge] < best_dist): 37 best_edge = edge 38 best_dist = edges[i][edge] 39 40 used.add(best_edge) 41 dist += best_dist 42 43 return dist