docs: union find

add size count

Author: QuBenhao

problems/problems_827/solution.py

@@ -22,7 +22,7 @@ def point_to_idx(x, y):
                     for dx, dy in DIRS:
                         if 0 <= (nx := i + dx) < n and 0 <= (ny := j + dy) < n and grid[nx][ny] == 1:
                             uf.union(p, point_to_idx(nx, ny))
-        ans = max(uf.count)
+        ans = max(uf.size)
         for i, row in enumerate(grid):
             for j, val in enumerate(row):
                 if val == 0:
@@ -34,26 +34,25 @@ def point_to_idx(x, y):
                             if root in explored:
                                 continue
                             explored.add(root)
-                            tot += uf.count[root]
+                            tot += uf.size[root]
                     ans = max(ans, tot)
 
         return ans
 
 class UnionFind:
-    def __init__(self, size):
-        self.parent = list(range(size))
-        self.rank = [1] * size
-        self.count = [1] * size
-        self.size = size
-        self.cc = size
+    def __init__(self, n: int):
+        self.parent = list(range(n))
+        self.rank = [1] * n
+        self.size = [1] * n
+        self.cc = n
 
-    def find(self, x):
+    def find(self, x: int) -> int:
         while self.parent[x] != x:
             self.parent[x] = self.parent[self.parent[x]]  # 路径压缩
             x = self.parent[x]
         return x
 
-    def union(self, x, y):
+    def union(self, x: int, y: int) -> bool:
         root_x = self.find(x)
         root_y = self.find(y)
 
@@ -63,14 +62,14 @@ def union(self, x, y):
         # 按秩合并
         if self.rank[root_x] > self.rank[root_y]:
             self.parent[root_y] = root_x
-            self.count[root_x] += self.count[root_y]
+            self.size[root_x] += self.size[root_y]
         else:
             self.parent[root_x] = root_y
             if self.rank[root_x] == self.rank[root_y]:
                 self.rank[root_y] += 1
-            self.count[root_y] += self.count[root_x]
+            self.size[root_y] += self.size[root_x]
         self.cc -= 1
         return True
 
-    def is_connected(self, x, y):
+    def is_connected(self, x:int, y:int) -> bool:
         return self.find(x) == self.find(y)

templates/data_structure/union_find_set.md

@@ -7,18 +7,19 @@
 
 ```python
 class UnionFind:
-    def __init__(self, size):
-        self.parent = list(range(size))
-        self.rank = [1] * size
-        self.size = size
+    def __init__(self, n: int):
+        self.parent = list(range(n))
+        self.rank = [1] * n
+        self.size = [1] * n
+        self.cc = n
 
-    def find(self, x):
+    def find(self, x: int) -> int:
         while self.parent[x] != x:
             self.parent[x] = self.parent[self.parent[x]]  # 路径压缩
             x = self.parent[x]
         return x
 
-    def union(self, x, y):
+    def union(self, x: int, y: int) -> bool:
         root_x = self.find(x)
         root_y = self.find(y)
 
@@ -28,15 +29,18 @@ class UnionFind:
         # 按秩合并
         if self.rank[root_x] > self.rank[root_y]:
             self.parent[root_y] = root_x
+            self.size[root_x] += self.size[root_y]
         else:
             self.parent[root_x] = root_y
             if self.rank[root_x] == self.rank[root_y]:
                 self.rank[root_y] += 1
-        self.size -= 1
+            self.size[root_y] += self.size[root_x]
+        self.cc -= 1
         return True
 
-    def is_connected(self, x, y):
+    def is_connected(self, x:int, y:int) -> bool:
         return self.find(x) == self.find(y)
+
 ```
 
 ```go