Optimize Water Distribution in a Village

lcpy gen -n 1168   # by problem number
lcpy gen -s optimize_water_distribution_in_a_village   # by problem name

Input: n = 3, wells = [1,2,2], pipes = [[1,2,1],[2,3,1]]
Output: 3
Explanation: The image shows the costs of connecting houses using pipes.
The best strategy is to build a well in the first house with cost 1 and connect the other houses to it with cost 2 so the total cost is 3.

Input: n = 2, wells = [1,1], pipes = [[1,2,1],[1,2,2]]
Output: 2
Explanation: We can supply water with cost two using one of the three options:
Option 1:
  - Build a well inside house 1 with cost 1.
  - Build a well inside house 2 with cost 1.
The total cost will be 2.
Option 2:
  - Build a well inside house 1 with cost 1.
  - Connect house 2 with house 1 with cost 1.
The total cost will be 2.
Option 3:
  - Build a well inside house 2 with cost 1.
  - Connect house 1 with house 2 with cost 1.
The total cost will be 2.
Note that we can connect houses 1 and 2 with cost 1 or with cost 2 but we will always choose **the cheapest option**.

class Solution:
    # Time: O((m + n) log(m + n)) where m = len(pipes), n = len(wells)
    # Space: O(n + m)
    def min_cost_to_supply_water(self, n: int, wells: list[int], pipes: list[list[int]]) -> int:
        # Virtual well node 0: connecting house i to it costs wells[i - 1].
        edges = [(w, 0, i + 1) for i, w in enumerate(wells)]
        edges += [(c, a, b) for a, b, c in pipes]
        edges.sort()

        parent = list(range(n + 1))

        def find(x: int) -> int:
            while parent[x] != x:
                parent[x] = parent[parent[x]]
                x = parent[x]
            return x

        total = 0
        components = n + 1
        for cost, a, b in edges:
            ra, rb = find(a), find(b)
            if ra == rb:
                continue
            parent[ra] = rb
            total += cost
            components -= 1
            if components == 1:
                break
        return total