Merge pull request #5 from harrymvr/develop

Adds option to stop adding edges to insertable

cliquetree/cliquetree.py

@@ -212,13 +212,18 @@ def add_edge(self, x, y, update_insertable=True):
         #     self.insertable.remove((x, y))
 
         self.uid += 1
+        self.insertable = set()
         if update_insertable:
-            self.insertable = set()
             for v in self.G:
                 self.update_insertable(v)
         return True
 
-    def update_insertable(self, v):
+    def update_insertable(self, v, stop_at=None):
+        """Updates the insertable edges in the graph.
+
+        For early stopping, set stop_at to k. Then, the function will return
+        after when k edges have been added to the insertable set.
+        """
         K1 = 0
         Kx = None
         cliques_visited = set()
@@ -255,6 +260,9 @@ def update_insertable(self, v):
                                     self.insertable.add(self._edge(u, v))
                                     if u == v:
                                         raise ValueError('u is equal to v')
+                                    if stop_at is not None and \
+                                            len(self.insertable) >= stop_at:
+                                        return
                     else:
                         min_weights.append(min_weights[-1])
                     if nodes_seen:
@@ -277,6 +285,9 @@ def update_insertable(self, v):
             if clq not in cliques_visited:
                 for u in self.nodes_in_clique[clq]:
                     self.insertable.add(self._edge(u, v))
+                    if stop_at is not None and \
+                            len(self.insertable) >= stop_at:
+                        return
 
     def update_deletable(self):
         self.deletable = set()

tests/test_cliquetree.py

@@ -33,5 +33,35 @@ def test_insertable1():
         assert c.insertable == insertable
 
 
+def test_update():
+    c = CliqueTree()
+    edges = [(1, 2), (2, 3), (3, 4), (4, 5), (5, 6), (6, 7), (3, 5), (6, 4),
+             (1, 3), (1, 5), (1, 4), (5, 7), (2, 5), (1, 6), (2, 4), (1, 7),
+             (4, 7), (2, 7), (3, 7), (2, 6), (3, 6)]
+    solutions = [frozenset([]), frozenset([(1, 3)]), frozenset([(1, 3), (2, 4)]),
+                 frozenset([(1, 3), (2, 4), (3, 5)]),
+                 frozenset([(4, 6), (1, 3), (2, 4), (3, 5)]),
+                 frozenset([(4, 6), (5, 7), (1, 3), (2, 4), (3, 5)]),
+                 frozenset([(1, 3), (4, 6), (5, 7), (3, 6), (2, 5), (2, 4)]),
+                 frozenset([(4, 7), (1, 3), (5, 7), (3, 6), (2, 5), (2, 4)]),
+                 frozenset([(4, 7), (5, 7), (1, 4), (1, 5), (3, 6), (2, 5), (2, 4)]),
+                 frozenset([(2, 5), (5, 7), (3, 6), (1, 4), (4, 7)]),
+                 frozenset([(4, 7), (5, 7), (1, 6), (3, 6), (2, 5), (2, 4)]),
+                 frozenset([(2, 5), (1, 6), (2, 4), (3, 6), (4, 7)]),
+                 frozenset([(4, 7), (1, 6), (2, 4), (3, 6)]),
+                 frozenset([(4, 7), (2, 4), (3, 6), (1, 7)]),
+                 frozenset([(4, 7), (2, 6), (3, 6), (1, 7)]),
+                 frozenset([(4, 7), (2, 6), (3, 6)]),
+                 frozenset([(2, 7), (3, 7), (2, 6), (3, 6)]),
+                 frozenset([(3, 7), (2, 6)]),
+                 frozenset([(2, 6), (3, 6)]),
+                 frozenset([(3, 6)]),
+                 frozenset([])]
 
-
+    for edge, insertable in zip(edges[:10], solutions[:10]):
+        c.add_edge(*edge)
+    c.add_edge(*edges[10], update_insertable=False)
+    assert len(c.insertable) == 0
+    c.update_insertable(7, stop_at=1)
+    assert len(c.insertable) == 1
+    assert len(c.insertable.intersection(solutions[10])) == 1