feat: add solutions to lc problem: No.1168 (#2052)

No.1168.Optimize Water Distribution in a Village

Author: yanglbme

solution/1100-1199/1168.Optimize Water Distribution in a Village/README.md

@@ -1,32 +1,56 @@
-class Solution {
-public:
-    int minCostToSupplyWater(int n, vector<int>& wells, vector<vector<int>>& pipes) {
-        for (int i = 0; i < n; ++i) {
-            pipes.push_back({0, i + 1, wells[i]});
-        }
-        sort(pipes.begin(), pipes.end(), [](const vector<int>& a, const vector<int>& b) {
-            return a[2] < b[2];
-        });
-        int p[n + 1];
-        iota(p, p + n + 1, 0);
-        function<int(int)> find = [&](int x) {
-            if (p[x] != x) {
-                p[x] = find(p[x]);
-            }
-            return p[x];
-        };
-        int ans = 0;
-        for (const auto& x : pipes) {
-            int pa = find(x[0]), pb = find(x[1]);
-            if (pa == pb) {
-                continue;
-            }
-            p[pa] = pb;
-            ans += x[2];
-            if (--n == 0) {
-                break;
-            }
-        }
-        return ans;
-    }
+class UnionFind {
+public:
+    UnionFind(int n) {
+        p = vector<int>(n);
+        size = vector<int>(n, 1);
+        iota(p.begin(), p.end(), 0);
+    }
+
+    bool unite(int a, int b) {
+        int pa = find(a), pb = find(b);
+        if (pa == pb) {
+            return false;
+        }
+        if (size[pa] > size[pb]) {
+            p[pb] = pa;
+            size[pa] += size[pb];
+        } else {
+            p[pa] = pb;
+            size[pb] += size[pa];
+        }
+        return true;
+    }
+
+    int find(int x) {
+        if (p[x] != x) {
+            p[x] = find(p[x]);
+        }
+        return p[x];
+    }
+
+private:
+    vector<int> p, size;
+};
+
+class Solution {
+public:
+    int minCostToSupplyWater(int n, vector<int>& wells, vector<vector<int>>& pipes) {
+        for (int i = 0; i < n; ++i) {
+            pipes.push_back({0, i + 1, wells[i]});
+        }
+        sort(pipes.begin(), pipes.end(), [](const vector<int>& a, const vector<int>& b) {
+            return a[2] < b[2];
+        });
+        UnionFind uf(n + 1);
+        int ans = 0;
+        for (const auto& x : pipes) {
+            if (uf.unite(x[0], x[1])) {
+                ans += x[2];
+                if (--n == 0) {
+                    break;
+                }
+            }
+        }
+        return ans;
+    }
 };

solution/1100-1199/1168.Optimize Water Distribution in a Village/README_EN.md

@@ -1,30 +1,52 @@
-func minCostToSupplyWater(n int, wells []int, pipes [][]int) (ans int) {
-	for i, w := range wells {
-		pipes = append(pipes, []int{0, i + 1, w})
-	}
-	sort.Slice(pipes, func(i, j int) bool { return pipes[i][2] < pipes[j][2] })
-	p := make([]int, n+1)
+type unionFind struct {
+	p, size []int
+}
+
+func newUnionFind(n int) *unionFind {
+	p := make([]int, n)
+	size := make([]int, n)
 	for i := range p {
 		p[i] = i
+		size[i] = 1
 	}
-	var find func(int) int
-	find = func(x int) int {
-		if p[x] != x {
-			p[x] = find(p[x])
-		}
-		return p[x]
+	return &unionFind{p, size}
+}
+
+func (uf *unionFind) find(x int) int {
+	if uf.p[x] != x {
+		uf.p[x] = uf.find(uf.p[x])
 	}
+	return uf.p[x]
+}
 
+func (uf *unionFind) union(a, b int) bool {
+	pa, pb := uf.find(a), uf.find(b)
+	if pa == pb {
+		return false
+	}
+	if uf.size[pa] > uf.size[pb] {
+		uf.p[pb] = pa
+		uf.size[pa] += uf.size[pb]
+	} else {
+		uf.p[pa] = pb
+		uf.size[pb] += uf.size[pa]
+	}
+	return true
+}
+
+func minCostToSupplyWater(n int, wells []int, pipes [][]int) (ans int) {
+	for i, w := range wells {
+		pipes = append(pipes, []int{0, i + 1, w})
+	}
+	sort.Slice(pipes, func(i, j int) bool { return pipes[i][2] < pipes[j][2] })
+	uf := newUnionFind(n + 1)
 	for _, x := range pipes {
-		pa, pb := find(x[0]), find(x[1])
-		if pa == pb {
-			continue
-		}
-		p[pa] = pb
-		ans += x[2]
-		n--
-		if n == 0 {
-			break
+		if uf.union(x[0], x[1]) {
+			ans += x[2]
+			n--
+			if n == 0 {
+				break
+			}
 		}
 	}
 	return

solution/1100-1199/1168.Optimize Water Distribution in a Village/Solution.cpp

@@ -1,39 +1,57 @@
-class Solution {
-    private int[] p;
-
-    public int minCostToSupplyWater(int n, int[] wells, int[][] pipes) {
-        var nums = new int[n + pipes.length][0];
-        int j = 0;
-        for (var pipe : pipes) {
-            nums[j++] = pipe;
-        }
-        for (int i = 0; i < n; ++i) {
-            nums[j++] = new int[] {0, i + 1, wells[i]};
-        }
-        Arrays.sort(nums, (a, b) -> a[2] - b[2]);
-        p = new int[n + 1];
-        for (int i = 1; i <= n; ++i) {
-            p[i] = i;
-        }
-        int ans = 0;
-        for (var x : nums) {
-            int pa = find(x[0]), pb = find(x[1]);
-            if (pa == pb) {
-                continue;
-            }
-            ans += x[2];
-            p[pa] = pb;
-            if (--n == 0) {
-                break;
-            }
-        }
-        return ans;
-    }
-
-    private int find(int x) {
-        if (p[x] != x) {
-            p[x] = find(p[x]);
-        }
-        return p[x];
-    }
+class UnionFind {
+    private int[] p;
+    private int[] size;
+
+    public UnionFind(int n) {
+        p = new int[n];
+        size = new int[n];
+        for (int i = 0; i < n; ++i) {
+            p[i] = i;
+            size[i] = 1;
+        }
+    }
+
+    public int find(int x) {
+        if (p[x] != x) {
+            p[x] = find(p[x]);
+        }
+        return p[x];
+    }
+
+    public boolean union(int a, int b) {
+        int pa = find(a), pb = find(b);
+        if (pa == pb) {
+            return false;
+        }
+        if (size[pa] > size[pb]) {
+            p[pb] = pa;
+            size[pa] += size[pb];
+        } else {
+            p[pa] = pb;
+            size[pb] += size[pa];
+        }
+        return true;
+    }
+}
+
+class Solution {
+    public int minCostToSupplyWater(int n, int[] wells, int[][] pipes) {
+        int[][] nums = Arrays.copyOf(pipes, pipes.length + n);
+        for (int i = 0; i < n; i++) {
+            nums[pipes.length + i] = new int[] {0, i + 1, wells[i]};
+        }
+        Arrays.sort(nums, (a, b) -> a[2] - b[2]);
+        UnionFind uf = new UnionFind(n + 1);
+        int ans = 0;
+        for (var x : nums) {
+            int a = x[0], b = x[1], c = x[2];
+            if (uf.union(a, b)) {
+                ans += c;
+                if (--n == 0) {
+                    break;
+                }
+            }
+        }
+        return ans;
+    }
 }

solution/1100-1199/1168.Optimize Water Distribution in a Village/Solution.go

@@ -1,24 +1,40 @@
-class Solution:
-    def minCostToSupplyWater(
-        self, n: int, wells: List[int], pipes: List[List[int]]
-    ) -> int:
-        def find(x: int) -> int:
-            if p[x] != x:
-                p[x] = find(p[x])
-            return p[x]
-
-        for i, w in enumerate(wells, 1):
-            pipes.append([0, i, w])
-        pipes.sort(key=lambda x: x[2])
-        p = list(range(n + 1))
-        ans = 0
-        for i, j, c in pipes:
-            pa, pb = find(i), find(j)
-            if pa == pb:
-                continue
-            p[pa] = pb
-            ans += c
-            n -= 1
-            if n == 0:
-                break
-        return ans
+class UnionFind:
+    __slots__ = ("p", "size")
+
+    def __init__(self, n):
+        self.p = list(range(n))
+        self.size = [1] * n
+
+    def find(self, x: int) -> int:
+        if self.p[x] != x:
+            self.p[x] = self.find(self.p[x])
+        return self.p[x]
+
+    def union(self, a: int, b: int) -> bool:
+        pa, pb = self.find(a), self.find(b)
+        if pa == pb:
+            return False
+        if self.size[pa] > self.size[pb]:
+            self.p[pb] = pa
+            self.size[pa] += self.size[pb]
+        else:
+            self.p[pa] = pb
+            self.size[pb] += self.size[pa]
+        return True
+
+
+class Solution:
+    def minCostToSupplyWater(
+        self, n: int, wells: List[int], pipes: List[List[int]]
+    ) -> int:
+        for i, w in enumerate(wells, 1):
+            pipes.append([0, i, w])
+        pipes.sort(key=lambda x: x[2])
+        uf = UnionFind(n + 1)
+        ans = 0
+        for a, b, c in pipes:
+            if uf.union(a, b):
+                ans += c
+                n -= 1
+                if n == 0:
+                    return ans

solution/1100-1199/1168.Optimize Water Distribution in a Village/Solution.java

@@ -0,0 +1,60 @@
+struct UnionFind {
+    p: Vec<usize>,
+    size: Vec<usize>,
+}
+
+impl UnionFind {
+    fn new(n: usize) -> Self {
+        let p: Vec<usize> = (0..n).collect();
+        let size = vec![1; n];
+        UnionFind { p, size }
+    }
+
+    fn find(&mut self, x: usize) -> usize {
+        if self.p[x] != x {
+            self.p[x] = self.find(self.p[x]);
+        }
+        self.p[x]
+    }
+
+    fn union(&mut self, a: usize, b: usize) -> bool {
+        let pa = self.find(a);
+        let pb = self.find(b);
+        if pa == pb {
+            false
+        } else if self.size[pa] > self.size[pb] {
+            self.p[pb] = pa;
+            self.size[pa] += self.size[pb];
+            true
+        } else {
+            self.p[pa] = pb;
+            self.size[pb] += self.size[pa];
+            true
+        }
+    }
+}
+
+impl Solution {
+    pub fn min_cost_to_supply_water(n: i32, wells: Vec<i32>, pipes: Vec<Vec<i32>>) -> i32 {
+        let n = n as usize;
+        let mut pipes = pipes;
+        for i in 0..n {
+            pipes.push(vec![0, (i + 1) as i32, wells[i]]);
+        }
+        pipes.sort_by(|a, b| a[2].cmp(&b[2]));
+        let mut uf = UnionFind::new(n + 1);
+        let mut ans = 0;
+        for pipe in pipes {
+            let a = pipe[0] as usize;
+            let b = pipe[1] as usize;
+            let c = pipe[2];
+            if uf.union(a, b) {
+                ans += c;
+                if n == 0 {
+                    break;
+                }
+            }
+        }
+        ans
+    }
+}