Leetcode 1168. Optimize Water Distribution in a Village

Problem Explanation

In this problem, we are given n number of houses in a city, a cost wells[i] to build a well in each house directly, and an array pipes[] indicating costs to install a bi-directional pipe between two houses.

The task is to calculate the minimum cost required to supply water to all houses. We can either directly build wells for each house or connect some houses with pipes and supply water from those wells.

This problem is similar to finding a minimum spanning tree in a graph where the vertices are houses and the distances are costs (wells or pipes).

Approach

We create a graph where each house is connected by a pipe (if provided) and also connect a dummy node at 0, with the well's cost as the edge's weight. The minimum cost for this problem is equivalent to finding the minimum cost to connect all vertices (houses), i.e., finding the Minimum Spanning Tree (MST) for this graph.

We can use Prim's algorithm to find the MST. In this algorithm, we first add the vertex with the minimum cost to a priority queue. Then we iterate until all vertices are not in the MST. In each iteration, we pop out the vertex with the smallest cost from the priority queue, add it to the MST, and update the queue with its connected vertices if they are not in MST already.

Python Solution

1
2python
3import heapq
4from collections import defaultdict
5
6class Solution:
7    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:
8        graph = defaultdict(list)
9        for u,v,cost in pipes:
10            graph[v].append((cost,u))
11            graph[u].append((cost, v))
12        for i in range(n):
13            graph[0].append((wells[i], i+1))
14        
15        mst = set([0])
16        minHeap = graph[0]
17
18        heapq.heapify(minHeap)
19        cost = 0
20        while len(minHeap)>0:
21            w,v =  heapq.heappop(minHeap)
22            if v not in mst:
23                mst.add(v)
24                cost += w
25                for ed in graph[v]:
26                    if ed[1] not in mst:
27                        heapq.heappush(minHeap, ed)
28        return cost

Java Solution

1
2java
3import java.util.*;
4
5class Solution {
6    public int minCostToSupplyWater(int n, int[] wells, int[][] pipes) {
7        List<List<int[]>> graph = new ArrayList<>();
8        for (int i = 0; i <= n; i++) {
9            graph.add(new ArrayList<>());
10        }
11        for (int[] pipe : pipes) {
12            int u = pipe[0], v = pipe[1], w = pipe[2];
13            graph.get(u).add(new int[]{v, w});
14            graph.get(v).add(new int[]{u, w});
15        }
16        for (int i = 0; i < n; i++) {
17            graph.get(0).add(new int[]{i + 1, wells[i]});
18        }
19
20        PriorityQueue<int[]> minHeap = new PriorityQueue<>(Comparator.comparingInt(o -> o[1]));
21        minHeap.offer(new int[]{0, 0});
22        boolean[] mstSet = new boolean[n + 1];
23        int totalCost = 0;
24
25        while (!minHeap.isEmpty()) {
26            int[] u = minHeap.poll();
27            if (mstSet[u[0]])
28                continue;
29            mstSet[u[0]] = true;
30            totalCost += u[1];
31            for (int[] v : graph.get(u[0])) {
32                if (!mstSet[v[0]])
33                    minHeap.offer(v);
34            }
35        }
36
37        return totalCost;
38    }
39}

JavaScript Solution

1
2javascript
3class MinHeap {
4    constructor() {
5        this.heap = [];
6    }
7
8    getMin() {
9        return this.heap[0];
10    }
11
12    isEmpty() {
13        return this.heap.length === 0;
14    }
15
16    insert(val) {
17        this.heap.push(val);
18        this.bubbleUp();
19    }
20
21    removeMin() {
22        if (this.heap.length === 1) {
23            return this.heap.pop();
24        } else {
25            let min = this.heap[0];
26            this.heap[0] = this.heap.pop();
27            this.bubbleDown();
28            return min;
29        }
30    }
31
32    bubbleDown() {
33        let index = 0;
34
35        while (this.leftChild(index) && 
36              (this.heap[index][0] > this.leftChild(index)[0] ||
37              (this.rightChild(index) && this.heap[index][0] > this.rightChild(index)[0]))) {
38            let smallerIndex = this.leftChildIndex(index);
39
40            if (this.rightChild(index) && this.rightChild(index)[0] < this.leftChild(index)[0]) {
41                smallerIndex = this.rightChildIndex(index);
42            }
43
44            this.swap(smallerIndex, index);
45            index = smallerIndex;
46        }
47    }
48
49    bubbleUp() {
50        let index = this.heap.length - 1;
51
52        while (this.parent(index) && this.parent(index)[0] > this.heap[index][0]) {
53            this.swap(this.parentIndex(index), index);
54            index = this.parentIndex(index);
55        }
56    }
57
58    parent(index) {
59        return this.heap[this.parentIndex(index)];
60    }
61
62    leftChild(index) {
63        return this.heap[this.leftChildIndex(index)];
64    }
65
66    rightChild(index) {
67        return this.heap[this.rightChildIndex(index)];
68    }
69
70    leftChildIndex(index) {
71        return (2 * index) + 1;
72    }
73
74    rightChildIndex(index) {
75        return (2 * index) + 2;
76    }
77
78    parentIndex(index) {
79        return Math.floor((index - 1) / 2);
80    }
81
82    swap(i, j) {
83        let temp = this.heap[i];
84        this.heap[i] = this.heap[j];
85        this.heap[j] = temp;
86    }
87}
88
89class Solution {
90    minCostToSupplyWater(n, wells, pipes) {
91        let graph = {};
92        pipes.forEach(pipe => {
93            if (!graph[pipe[0]]) {
94                graph[pipe[0]] = {};
95            }
96            graph[pipe[0]][pipe[1]] = pipe[2];
97            if (!graph[pipe[1]]) {
98                graph[pipe[1]] = {};
99            }
100            graph[pipe[1]][pipe[0]] = pipe[2];
101        });
102        
103        wells.forEach((well, idx) => {
104            if (!graph[0]) {
105                graph[0] = {};
106            }
107            graph[0][idx+1] = well;
108        });
109        
110        let visited = {}, totalCost = 0;
111        let minHeap = new MinHeap();
112        Object.keys(graph[0]).forEach(i => minHeap.insert([graph[0][i], i]));
113    
114        while (!minHeap.isEmpty()) {
115            let minCost = minHeap.removeMin();
116            if (visited[minCost[1]])
117                continue;
118            visited[minCost[1]] = true;
119            totalCost += minCost[0];
120            Object.keys(graph[minCost[1]]).forEach(house => {
121                if (!visited[house])
122                    minHeap.insert([graph[minCost[1]][house], house]);
123            });
124        }
125        return totalCost;
126    }
127}

The above three solutions are quite detailed and efficient to solve the problem. However, there're some important takeaways, regardless of the language we work with:

  1. Use Prim's algorithm to minimum spanning tree: It's a greedy algorithm that strives to find the minimum cost of connecting all nodes in the graph. It does so by adding the least costly node (that's not yet added to the MST) at each iteration.

  2. Using a min-heap: A min-heap helps to get the node with the minimum cost in O(logn) time, which is a significant improvement over a simple linear search.

  3. Marking nodes as 'visited': After a node is added to the MST, mark it as visited. No further consideration is necessary for the node because we're looking for an acyclic connected graph - a tree.

  4. Connecting all houses to a dummy node at 0: It is an ingenious way to add the well costs as edges in the graph. This node is the starting point of our algorithm. Without adding this node, the well costs wouldn't have been considered by prim's algorithm.

Remember, the priority is to understand the problem's essence and solution. Be it Python, JavaScript, or Java, the fundamental concept remains the same.


Got a question? Ask the Teaching Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.


TA 👨‍🏫