From 0693252f1fbdd4b2fc3db010b25305cd55eb1406 Mon Sep 17 00:00:00 2001
From: Stefan Huber <shuber@sthu.org>
Date: Wed, 23 May 2018 19:03:51 +0200
Subject: [PATCH] union_find: Do not cummulate weight

---
 union_find.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/union_find.py b/union_find.py
index 320bba1..0cd8604 100644
--- a/union_find.py
+++ b/union_find.py
@@ -29,11 +29,20 @@ class UnionFind:
         self.weights = {}
         self.parents = {}
 
+    def add(self, object, weight):
+        if object not in self.parents:
+            self.parents[object] = object
+            self.weights[object] = weight
+
+    def __contains__(self, object):
+        return object in self.parents
+
     def __getitem__(self, object):
         """Find and return the name of the set containing the object."""
 
         # check for previously unknown object
         if object not in self.parents:
+            assert(False)
             self.parents[object] = object
             self.weights[object] = 1
             return object
@@ -62,5 +71,4 @@ class UnionFind:
         heaviest = max([(self.weights[r], r) for r in roots])[1]
         for r in roots:
             if r != heaviest:
-                self.weights[heaviest] += self.weights[r]
                 self.parents[r] = heaviest
-- 
2.39.5