Min Cost to Connect All Points LeetCode solution in Python

1
# Runtime: 2922 ms (Top 60.17%) | Memory: 79.1 MB (Top 76.18%)
2
class Solution:
3
    def minCostConnectPoints(self, points: List[List[int]]) -> int:
4

5
        cost = 0
6
        heap = []
7

8
        # set to store past points to prevent cycle
9
        visited = set([0])
10

11
        # i == the index of current point
12
        i = 0
13

14
        while len(visited) < len(points):
15
            # Add all costs from current point to unreached points to min heap
16
            for j in range(len(points)):
17
                if j == i or j in visited:
18
                    continue
19
                distance = abs(points[i][0] - points[j][0]) + abs(
20
                    points[i][1] - points[j][1]
21
                )
22
                heapq.heappush(heap, (distance, j))
23

24
            # Add new min cost edge
25
            while True:
26
                dist, point = heapq.heappop(heap)
27
                if point not in visited:
28
                    cost += dist
29
                    # Add point to visited to prevent cycle
30
                    visited.add(point)
31
                    # Update point
32
                    i = point
33
                    break
34

35
        return cost