Water Plant Infrastructure Optimization

Problem

A group of houses is connected to the main water plant by means of a set of pipes. A house can either be connected by a set of pipes extending directly to the plant, or indirectly by a pipe to a nearby house which is otherwise connected.

For example, here is a possible configuration, where A, B, and C are houses, and arrows represent pipes:

1

A <--> B <--> C <--> plant

Each pipe has an associated cost, which the utility company would like to minimize. Given an undirected graph of pipe connections, return the lowest cost configuration of pipes such that each house has access to water.

Note: This problem is similar to Connecting Cities With Minimum Cost but focuses on water distribution infrastructure with adjacency list input format.

Examples

Example 1

1
2
3


Input: pipes = {'plant': {'A': 1, 'B': 5, 'C': 20}, 'A': {'C': 15}, 'B': {'C': 10}, 'C': {}}
Output: 16
Explanation: We can connect plant->A (cost 1), plant->B (cost 5), and B->C (cost 10) for total cost 16. This ensures all houses A, B, C have access to water through the minimum spanning tree.

Example 2

1
2
3


Input: pipes = {'plant': {'A': 10, 'B': 15}, 'A': {'B': 5}, 'B': {}}
Output: 15
Explanation: We can connect plant->A (cost 10) and A->B (cost 5) for total cost 15. Alternative: plant->B (cost 15) and plant->A (cost 10) for total cost 25, but first option is cheaper.

Example 3

1
2
3


Input: pipes = {'plant': {'A': 8}, 'A': {}}
Output: 8
Explanation: Only one house A needs to be connected directly to plant with cost 8.

Example 4

1
2
3


Input: pipes = {'plant': {}}
Output: 0
Explanation: No houses to connect, so cost is 0.

Example 5

1
2
3


Input: pipes = {'plant': {'A': 3, 'B': 7}, 'A': {}, 'B': {}}
Output: 10
Explanation: Both houses A and B must be connected directly to plant since there are no inter-house connections. Total cost = 3 + 7 = 10.

Solution

Method 1 - Kruskal’s Algorithm with Union-Find

Intuition

This is a Minimum Spanning Tree (MST) problem where we need to connect all houses to the water plant with minimum cost. We can use Kruskal’s algorithm by first converting the adjacency list to an edge list, then applying the standard MST approach with Union-Find to avoid cycles.

Approach

Convert the adjacency list representation to a list of edges with costs
Sort all edges by cost in ascending order
Initialize Union-Find data structure for all nodes (including plant)
For each edge in sorted order, if it connects different components, add it to MST
Continue until all houses are connected to the plant component
Return the total cost of selected edges

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50


class Solution {
public:
    int minCostWaterPlant(unordered_map<string, unordered_map<string, int>>& pipes) {
        if (pipes.empty() || pipes["plant"].empty()) return 0;
        
        // Convert adjacency list to edge list
        vector<tuple<int, string, string>> edges;
        unordered_set<string> nodes;
        
        for (auto& [from, neighbors] : pipes) {
            nodes.insert(from);
            for (auto& [to, cost] : neighbors) {
                nodes.insert(to);
                edges.push_back({cost, from, to});
            }
        }
        
        // Sort edges by cost
        sort(edges.begin(), edges.end());
        
        // Union-Find
        unordered_map<string, string> parent;
        for (const string& node : nodes) {
            parent[node] = node;
        }
        
        function<string(const string&)> find = [&](const string& x) -> string {
            return parent[x] == x ? x : parent[x] = find(parent[x]);
        };
        
        int totalCost = 0;
        int edgesUsed = 0;
        int totalNodes = nodes.size();
        
        for (auto& [cost, u, v] : edges) {
            string rootU = find(u);
            string rootV = find(v);
            
            if (rootU != rootV) {
                parent[rootU] = rootV;
                totalCost += cost;
                edgesUsed++;
                
                if (edgesUsed == totalNodes - 1) break;
            }
        }
        
        return totalCost;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62


func minCostWaterPlant(pipes map[string]map[string]int) int {
    if len(pipes) == 0 || len(pipes["plant"]) == 0 {
        return 0
    }
    
    // Convert adjacency list to edge list
    type Edge struct {
        cost int
        from, to string
    }
    
    var edges []Edge
    nodes := make(map[string]bool)
    
    for from, neighbors := range pipes {
        nodes[from] = true
        for to, cost := range neighbors {
            nodes[to] = true
            edges = append(edges, Edge{cost, from, to})
        }
    }
    
    // Sort edges by cost
    sort.Slice(edges, func(i, j int) bool {
        return edges[i].cost < edges[j].cost
    })
    
    // Union-Find
    parent := make(map[string]string)
    for node := range nodes {
        parent[node] = node
    }
    
    var find func(string) string
    find = func(x string) string {
        if parent[x] != x {
            parent[x] = find(parent[x])
        }
        return parent[x]
    }
    
    totalCost := 0
    edgesUsed := 0
    totalNodes := len(nodes)
    
    for _, edge := range edges {
        rootU := find(edge.from)
        rootV := find(edge.to)
        
        if rootU != rootV {
            parent[rootU] = rootV
            totalCost += edge.cost
            edgesUsed++
            
            if edgesUsed == totalNodes-1 {
                break
            }
        }
    }
    
    return totalCost
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67


class Solution {
    public int minCostWaterPlant(Map<String, Map<String, Integer>> pipes) {
        if (pipes.isEmpty() || pipes.get("plant").isEmpty()) return 0;
        
        // Convert adjacency list to edge list
        List<int[]> edges = new ArrayList<>();
        Set<String> nodes = new HashSet<>();
        Map<String, Integer> nodeIndex = new HashMap<>();
        
        for (Map.Entry<String, Map<String, Integer>> entry : pipes.entrySet()) {
            String from = entry.getKey();
            nodes.add(from);
            for (Map.Entry<String, Integer> neighbor : entry.getValue().entrySet()) {
                String to = neighbor.getKey();
                int cost = neighbor.getValue();
                nodes.add(to);
            }
        }
        
        // Create node index mapping
        int idx = 0;
        for (String node : nodes) {
            nodeIndex.put(node, idx++);
        }
        
        // Build edges with indices
        for (Map.Entry<String, Map<String, Integer>> entry : pipes.entrySet()) {
            String from = entry.getKey();
            for (Map.Entry<String, Integer> neighbor : entry.getValue().entrySet()) {
                String to = neighbor.getKey();
                int cost = neighbor.getValue();
                edges.add(new int[]{cost, nodeIndex.get(from), nodeIndex.get(to)});
            }
        }
        
        // Sort edges by cost
        edges.sort((a, b) -> Integer.compare(a[0], b[0]));
        
        // Union-Find
        int n = nodes.size();
        int[] parent = new int[n];
        for (int i = 0; i < n; i++) parent[i] = i;
        
        int totalCost = 0;
        int edgesUsed = 0;
        
        for (int[] edge : edges) {
            int cost = edge[0], u = edge[1], v = edge[2];
            int rootU = find(parent, u);
            int rootV = find(parent, v);
            
            if (rootU != rootV) {
                parent[rootU] = rootV;
                totalCost += cost;
                edgesUsed++;
                
                if (edgesUsed == n - 1) break;
            }
        }
        
        return totalCost;
    }
    
    private int find(int[] parent, int x) {
        return parent[x] == x ? x : (parent[x] = find(parent, parent[x]));
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43


class Solution:
    def min_cost_water_plant(self, pipes: Dict[str, Dict[str, int]]) -> int:
        if not pipes or not pipes.get("plant", {}):
            return 0
        
        # Convert adjacency list to edge list
        edges = []
        nodes = set()
        
        for from_node, neighbors in pipes.items():
            nodes.add(from_node)
            for to_node, cost in neighbors.items():
                nodes.add(to_node)
                edges.append((cost, from_node, to_node))
        
        # Sort edges by cost
        edges.sort()
        
        # Union-Find
        parent = {node: node for node in nodes}
        
        def find(x: str) -> str:
            if parent[x] != x:
                parent[x] = find(parent[x])
            return parent[x]
        
        total_cost = 0
        edges_used = 0
        total_nodes = len(nodes)
        
        for cost, u, v in edges:
            root_u = find(u)
            root_v = find(v)
            
            if root_u != root_v:
                parent[root_u] = root_v
                total_cost += cost
                edges_used += 1
                
                if edges_used == total_nodes - 1:
                    break
        
        return total_cost

Complexity

⏰ Time complexity: O(E log E + E α(V)), where E is the number of edges and V is the number of nodes. Sorting edges takes O(E log E), and Union-Find operations take nearly constant time
🧺 Space complexity: O(V + E), for storing the edge list and Union-Find parent array

Method 2 - Prim’s Algorithm with Priority Queue

Intuition

Prim’s algorithm builds the MST by starting from any node (in our case, the plant) and greedily adding the minimum cost edge that connects a new node to the existing MST. This approach works directly with the adjacency list representation without conversion.

Approach

Start from the “plant” node as the root of MST
Use a priority queue to track edges from MST to non-MST nodes
Initialize with all edges from plant to its neighbors
Repeatedly extract minimum cost edge that connects to a new node
Add the new node to MST and push its edges to the queue
Continue until all nodes are in MST
Return total cost of edges used

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39


class Solution {
public:
    int minCostWaterPlant(unordered_map<string, unordered_map<string, int>>& pipes) {
        if (pipes.empty() || pipes.find("plant") == pipes.end()) return 0;
        
        unordered_set<string> inMST;
        priority_queue<pair<int, string>, vector<pair<int, string>>, greater<>> pq;
        
        // Start from plant
        inMST.insert("plant");
        for (auto& [neighbor, cost] : pipes["plant"]) {
            pq.push({cost, neighbor});
        }
        
        int totalCost = 0;
        
        while (!pq.empty()) {
            auto [cost, node] = pq.top();
            pq.pop();
            
            if (inMST.count(node)) continue;
            
            // Add node to MST
            inMST.insert(node);
            totalCost += cost;
            
            // Add edges from new node to queue
            if (pipes.count(node)) {
                for (auto& [neighbor, edgeCost] : pipes[node]) {
                    if (!inMST.count(neighbor)) {
                        pq.push({edgeCost, neighbor});
                    }
                }
            }
        }
        
        return totalCost;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61


func minCostWaterPlantPrim(pipes map[string]map[string]int) int {
    if len(pipes) == 0 || pipes["plant"] == nil {
        return 0
    }
    
    inMST := make(map[string]bool)
    
    // Priority queue implementation
    type Edge struct {
        cost int
        node string
    }
    
    pq := &PriorityQueue{}
    heap.Init(pq)
    
    // Start from plant
    inMST["plant"] = true
    for neighbor, cost := range pipes["plant"] {
        heap.Push(pq, Edge{cost, neighbor})
    }
    
    totalCost := 0
    
    for pq.Len() > 0 {
        edge := heap.Pop(pq).(Edge)
        
        if inMST[edge.node] {
            continue
        }
        
        // Add node to MST
        inMST[edge.node] = true
        totalCost += edge.cost
        
        // Add edges from new node to queue
        if neighbors, exists := pipes[edge.node]; exists {
            for neighbor, cost := range neighbors {
                if !inMST[neighbor] {
                    heap.Push(pq, Edge{cost, neighbor})
                }
            }
        }
    }
    
    return totalCost
}

type PriorityQueue []Edge

func (pq PriorityQueue) Len() int { return len(pq) }
func (pq PriorityQueue) Less(i, j int) bool { return pq[i].cost < pq[j].cost }
func (pq PriorityQueue) Swap(i, j int) { pq[i], pq[j] = pq[j], pq[i] }
func (pq *PriorityQueue) Push(x interface{}) { *pq = append(*pq, x.(Edge)) }
func (pq *PriorityQueue) Pop() interface{} {
    old := *pq
    n := len(old)
    item := old[n-1]
    *pq = old[0 : n-1]
    return item
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59


class Solution {
    public int minCostWaterPlant(Map<String, Map<String, Integer>> pipes) {
        if (pipes.isEmpty() || !pipes.containsKey("plant")) return 0;
        
        Set<String> inMST = new HashSet<>();
        PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) -> Integer.compare(a[0], b[0]));
        Map<String, Integer> nodeIndex = new HashMap<>();
        
        // Create node index mapping
        int idx = 0;
        for (String node : pipes.keySet()) {
            nodeIndex.put(node, idx++);
            for (String neighbor : pipes.get(node).keySet()) {
                if (!nodeIndex.containsKey(neighbor)) {
                    nodeIndex.put(neighbor, idx++);
                }
            }
        }
        
        // Start from plant
        inMST.add("plant");
        for (Map.Entry<String, Integer> entry : pipes.get("plant").entrySet()) {
            String neighbor = entry.getKey();
            int cost = entry.getValue();
            pq.offer(new int[]{cost, nodeIndex.get(neighbor)});
        }
        
        int totalCost = 0;
        String[] indexToNode = new String[nodeIndex.size()];
        for (Map.Entry<String, Integer> entry : nodeIndex.entrySet()) {
            indexToNode[entry.getValue()] = entry.getKey();
        }
        
        while (!pq.isEmpty()) {
            int[] edge = pq.poll();
            int cost = edge[0];
            String node = indexToNode[edge[1]];
            
            if (inMST.contains(node)) continue;
            
            // Add node to MST
            inMST.add(node);
            totalCost += cost;
            
            // Add edges from new node to queue
            if (pipes.containsKey(node)) {
                for (Map.Entry<String, Integer> entry : pipes.get(node).entrySet()) {
                    String neighbor = entry.getKey();
                    int edgeCost = entry.getValue();
                    if (!inMST.contains(neighbor)) {
                        pq.offer(new int[]{edgeCost, nodeIndex.get(neighbor)});
                    }
                }
            }
        }
        
        return totalCost;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34


class Solution:
    def min_cost_water_plant(self, pipes: Dict[str, Dict[str, int]]) -> int:
        if not pipes or "plant" not in pipes:
            return 0
        
        import heapq
        
        in_mst = set()
        pq = []
        
        # Start from plant
        in_mst.add("plant")
        for neighbor, cost in pipes["plant"].items():
            heapq.heappush(pq, (cost, neighbor))
        
        total_cost = 0
        
        while pq:
            cost, node = heapq.heappop(pq)
            
            if node in in_mst:
                continue
            
            # Add node to MST
            in_mst.add(node)
            total_cost += cost
            
            # Add edges from new node to queue
            if node in pipes:
                for neighbor, edge_cost in pipes[node].items():
                    if neighbor not in in_mst:
                        heapq.heappush(pq, (edge_cost, neighbor))
        
        return total_cost

Complexity

⏰ Time complexity: O(E log V), where E is the number of edges and V is the number of nodes. Each edge is processed at most once and priority queue operations take O(log V)
🧺 Space complexity: O(V + E), for the MST tracking set and priority queue storing edges

Problem#

Examples#

Example 1#

Example 2#

Example 3#

Example 4#

Example 5#

Solution#

Method 1 - Kruskal’s Algorithm with Union-Find#

Intuition#

Approach#

Code#

Complexity#

Method 2 - Prim’s Algorithm with Priority Queue#

Intuition#

Approach#

Code#

Complexity#

Problem

Examples

Example 1

Example 2

Example 3

Example 4

Example 5

Solution

Method 1 - Kruskal’s Algorithm with Union-Find

Intuition

Approach

Code

Complexity

Method 2 - Prim’s Algorithm with Priority Queue

Intuition

Approach

Code

Complexity