69 lines
2.5 KiB
Python
69 lines
2.5 KiB
Python
# Time: O(N^2·logN)
|
|
# Space: O(N^2) ; we could be pushing N·N(N - 1) / 2 edges into the heap
|
|
from collections import defaultdict
|
|
from heapq import heappush, heappop
|
|
|
|
class Solution:
|
|
def minCostConnectPoints(self, points: List[List[int]]) -> int:
|
|
'''
|
|
BFS + Prim's algorithm
|
|
'''
|
|
N = len(points)
|
|
neighbors = defaultdict(list)
|
|
|
|
# Build adjacency lists by connecting each point to every other
|
|
# point since we are building a new graph here
|
|
for i in range(N):
|
|
x1, y1 = points[i]
|
|
|
|
for j in range(i + 1, N):
|
|
x2, y2 = points[j]
|
|
|
|
# We are considering the Manhattan Distance as stated in
|
|
# the problem
|
|
distance = abs(x1 - x2) + abs(y1 - y2)
|
|
|
|
# We are going to add both the neighbor of given point as
|
|
# well as the distance so that we can later use it to create
|
|
# Prim's min heap
|
|
neighbors[i].append((distance, j))
|
|
neighbors[j].append((distance, i))
|
|
|
|
|
|
# Perform Prim's algorithm
|
|
|
|
# To avoid cycles
|
|
visited = set()
|
|
|
|
# We start at the first point, distance would be 0 since it's the first
|
|
# point
|
|
min_heap = [(0, 0)]
|
|
|
|
# We need to find distances from one node to all other nodes (as long as
|
|
# its not already visited), keep track of it in the min heap and once
|
|
# all neighbors of a node is visited, we should have found the next node
|
|
# to pop from the heap since it'd have the least distance.
|
|
#
|
|
# We need to keep doing this until we have visited all the nodes
|
|
total = 0
|
|
|
|
while len(visited) < N:
|
|
curr_dist, min_point = heappop(min_heap)
|
|
|
|
# It's possible we've already visited the node (let's say the very first node we
|
|
# initialized the heap with) since it could end up being added to the heap again
|
|
# while visiting other nodes. If this is the case, we skip this iteration of the
|
|
# loop
|
|
if min_point in visited:
|
|
continue
|
|
|
|
total += curr_dist
|
|
visited.add(min_point)
|
|
|
|
for neighbor in neighbors[min_point]:
|
|
if neighbor[1] not in visited:
|
|
heappush(min_heap, neighbor)
|
|
|
|
return total
|
|
|
|
|