Leetcode 1168 Optimize Water Distribution in a Village

Problem statement

https://leetcode.com/problems/optimize-water-distribution-in-a-village/

Solution

We need to notice, that if we are not allowed to use wells, this problem is nothing else but minimum spanning tree problem. Now, we can add node with index 0 and connect it with all other nodes with weights equal to cost of wells.

Complexity

Time complexity is O((E + n)*log E), space is O(E + n).

Code

class DSU:
    def __init__(self, N):
        self.p = list(range(N))

    def find(self, x):
        if self.p[x] != x:
            self.p[x] = self.find(self.p[x])
        return self.p[x]

    def union(self, x, y):
        xr = self.find(x)
        yr = self.find(y)
        self.p[xr] = yr

class Solution:
    def minCostToSupplyWater(self, n, wells, pipes):
        graph_wells = [[cost, 0, i] for i, cost in enumerate(wells, 1)]
        graph_pipes = [[cost, i, j] for i, j, cost in pipes]
        ans = 0

        dsu = DSU(n+1)

        for w, x, y in sorted(graph_wells + graph_pipes):
            if dsu.find(x) == dsu.find(y): continue
            dsu.union(x, y)
            ans += w
        
        return ans