Skip to content

1168. Optimize Water Distribution in a Village 👍

  • Time: $O((|V| + |E|)\log (|V| + |E|))$
  • Space: $O(|V| + |E|)$
 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
class Solution {
 public:
  int minCostToSupplyWater(int n, vector<int>& wells,
                           vector<vector<int>>& pipes) {
    int ans = 0;
    using P = pair<int, int>;
    vector<vector<P>> graph(n + 1);
    priority_queue<P, vector<P>, greater<>> minHeap;  // (d, u)

    for (const vector<int>& pipe : pipes) {
      const int u = pipe[0];
      const int v = pipe[1];
      const int w = pipe[2];
      graph[u].emplace_back(v, w);
      graph[v].emplace_back(u, w);
    }

    // Connect virtual 0 with nodes 1 to n.
    for (int i = 0; i < n; ++i) {
      graph[0].emplace_back(i + 1, wells[i]);
      minHeap.emplace(wells[i], i + 1);
    }

    unordered_set<int> mst{{0}};

    while (mst.size() < n + 1) {
      const auto [d, u] = minHeap.top();
      minHeap.pop();
      if (mst.contains(u))
        continue;
      // Add the new vertex.
      mst.insert(u);
      ans += d;
      // Expand if possible.
      for (const auto [v, w] : graph[u])
        if (!mst.contains(v))
          minHeap.emplace(w, v);
    }

    return ans;
  }
};

 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
class Solution {
  public int minCostToSupplyWater(int n, int[] wells, int[][] pipes) {
    int ans = 0;
    List<Pair<Integer, Integer>>[] graph = new List[n + 1];
    Queue<Pair<Integer, Integer>> minHeap =
        new PriorityQueue<>(Comparator.comparing(Pair::getKey)); // (d, u)

    for (int i = 0; i <= n; ++i)
      graph[i] = new ArrayList<>();

    for (int[] pipe : pipes) {
      final int u = pipe[0];
      final int v = pipe[1];
      final int w = pipe[2];
      graph[u].add(new Pair<>(v, w));
      graph[v].add(new Pair<>(u, w));
    }

    // Connect virtual 0 with nodes 1 to n.
    for (int i = 0; i < n; ++i) {
      graph[0].add(new Pair<>(i + 1, wells[i]));
      minHeap.offer(new Pair<>(wells[i], i + 1));
    }

    Set<Integer> mst = new HashSet<>(Arrays.asList(0));

    while (mst.size() < n + 1) {
      Pair<Integer, Integer> pair = minHeap.poll();
      final int d = pair.getKey();
      final int u = pair.getValue();
      if (mst.contains(u))
        continue;
      // Add the new vertex.
      mst.add(u);
      ans += d;
      // Expand if possible.
      for (Pair<Integer, Integer> pair : graph[u]) {
        final int v = pair.getKey();
        final int w = pair.getValue();
        if (!mst.contains(v))
          minHeap.offer(new Pair<>(w, v));
      }
    }

    return ans;
  }
}

 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
class Solution:
  def minCostToSupplyWater(
      self,
      n: int,
      wells: list[int],
      pipes: list[list[int]],
  ) -> int:
    ans = 0
    graph = [[] for _ in range(n + 1)]
    minHeap = []  # (d, u)

    for u, v, w in pipes:
      graph[u].append((v, w))
      graph[v].append((u, w))

    # Connect virtual 0 with nodes 1 to n.
    for i, well in enumerate(wells):
      graph[0].append((i + 1, well))
      heapq.heappush(minHeap, (well, i + 1))

    mst = {0}

    while len(mst) < n + 1:
      d, u = heapq.heappop(minHeap)
      if u in mst:
        continue
      # Add the new vertex.
      mst.add(u)
      ans += d
      # Expand if possible.
      for v, w in graph[u]:
        if v not in mst:
          heapq.heappush(minHeap, (w, v))

    return ans