import heapq
# This function is used to initialize the DSU.
def make(i, parent, sizez):
parent[i] = i
sizez[i] = 1
# This function is used to find the parent of each component of DSU
def find(v, parent):
if v == parent[v]:
return v
parent[v] = find(parent[v], parent)
return parent[v]
# This function is used to merge two components of DSU
def union(a, b, parent, sizez):
a = find(a, parent)
b = find(b, parent)
if a != b:
if sizez[a] < sizez[b]:
a, b = b, a
parent[b] = a
sizez[a] += sizez[b]
# Function to solve the problem
def min_cost_to_provide_water(n, wells, pipes):
edges = []
# Inserting the edges into the heap sorted by weights
for pipe in pipes:
u, v, wt = pipe
edges.append((wt, u, v))
# Inserting the edges created by the pseudo-vertex i.e., 0
for i in range(n):
edges.append((wells[i], 0, i + 1))
# Initializing DSU
parent = [0] * (n + 1)
sizez = [0] * (n + 1)
for i in range(n + 1):
make(i, parent, sizez)
answer = 0
# Applying Kruskal's MST algorithm
edges.sort()
for edge in edges:
wt, u, v = edge
if find(u, parent) != find(v, parent):
answer += wt
union(u, v, parent, sizez)
return answer
# Driver Function
if __name__ == "__main__":
N = 3
wells = [1, 2, 2]
pipes = [
(1, 2, 1),
(2, 3, 1)
]
# Function call
print(min_cost_to_provide_water(N, wells, pipes))