Merge remote-tracking branch 'origin/union-find'

1 files changed, 64 insertions, 2 deletions

diff --git a/store.py b/store.py
index 315b836..c988e47 100644
--- a/store.py
+++ b/store.py
@@ -9,6 +9,35 @@ class Node:
         self.lng = float(lng)
         self.neighbours = []
 
+        self.parent = None
+        self.size = 1
+
+    def find_root(self):
+        """ Recursively find the root of this node. """
+        if not self.parent:
+            return self
+        else:
+            return self.parent.find_root()
+
+    def union(self, other):
+        """ Mark this node and some other node as part of the same tree. """
+        this = self.find_root()
+        other = other.find_root()
+
+        # If we share the same root we are already part of the same tree
+        if this == other:
+            return
+
+        # Minimize tree size. This ensures optimal complexity.
+        # We could also store the rank instead.
+        if this.size < other.size:
+            this, other = other, this
+
+        # This is safe since other is a root node and currently doesn't have a
+        # parent.
+        other.parent = this
+        this.size += other.size
+
     def coord_tuple(self):
         return self.lat, self.lng
 
@@ -29,6 +58,8 @@ def add_neighbours(nodes):
 
             nodes[node1].neighbours.append(nodes[node2])
             nodes[node2].neighbours.append(nodes[node1])
+            # These two are neighbours and should be part of the same tree.
+            nodes[node1].union(nodes[node2])
 
     return nodes
 
@@ -45,11 +76,42 @@ def extract_osm_nodes(f_name):
 
     add_neighbours(nodes)
 
-    # remove nodes without neighbours.
+    # Remove nodes without neighbours
     for node_id, node in nodes.copy().items():
         if not node.neighbours:
             del nodes[node_id]
 
+    # Construct a forest of disjoint trees.
+    # The forest is a { root: tree }-dict. If we construct a relation where
+    # two nodes relate to each other if they are a part of the same tree,
+    # this can be shown to be an equivalence relation where the equivalence
+    # class is some node in the tree. In our case, the equivalence class
+    # becomes the root of each node since it will be the same for all nodes
+    # in the same tree due our use of union-find.
+    forest = defaultdict(dict)
+    for node_id, node in nodes.items():
+        forest[node.find_root().id][node_id] = node
+
+    # Find the largest disjoint tree.
+    # We store the root so we can easily test if a node is part of this
+    # tree or not.
+    best_size = 0
+    best_tree = None
+    best_root = None
+    for root in forest:
+        tree = forest[root]
+        size = len(tree)
+        if size > best_size:
+            best_size = size
+            best_tree = tree
+            best_root = root
+
+    unconnected_nodes = []
+    for _, node in nodes.items():
+        if node.find_root().id != best_root:
+            unconnected_nodes.append(node)
+
+    nodes = best_tree
     # create a "grid" by grouping nearby nodes.
     for node in nodes.copy().values():
 
@@ -62,7 +124,7 @@ def extract_osm_nodes(f_name):
         else:
             grid[key] = [node]
 
-    return nodes, grid
+    return nodes, grid, unconnected_nodes
 
 
 def select_nodes_in_rectangle(nodes, min_lat, max_lat, min_long, max_long):