mythical-programmer · June 22, 2017 17:58 · Jun 22, 2017
diff --git a/uf.py b/uf.py
@@ -0,0 +1,52 @@
+# union find data structure
+# general for any key (not just integer elements)
+# $ py uf.py in/tinyUF.txt 
+import sys
+from collections import defaultdict
+
+# weighted union-find data structure with path compression
+class UF:
+    def __init__(self):
+        self._sz = defaultdict(int)
+        self._id = {}
+
+    def _root(self, x):
+        # if not found, initialize element as it's own set
+        if x not in self._id:
+            self._id[x] = x
+        while x != self._id[x]:
+            self._id[x] = self._id[self._id[x]] # update to grandparent
+            x = self._id[x]
+        return x
+
+    def find(self, p, q):
+        return self._root(p) == self._root(q)
+
+    def union(self, p, q):
+        x = self._root(p)
+        y = self._root(q)
+        if x == y:
+            return
+        if self._sz[x] < self._sz[y]:
+            self._id[x] = y
+            self._sz[y] += self._sz[x]
+        else:
+            self._id[y] = x
+            self._sz[x] += self._sz[y]
+
+    def __str__(self):
+        return str(self._id)
+
+def main():
+    uf = UF()
+    filename = sys.argv[1]
+    with open(filename, 'r') as f:
+        line = f.readline().strip()
+        while line:
+            v, w = line.strip().split()
+            uf.union(v, w)
+            line = f.readline().strip()
+    print(uf)
+
+if __name__ == '__main__':
+    main()