groupConnectivityNaive: counting of satisfying flows

groupConnectivityNaive.py

@@ -30,14 +30,19 @@ def flowEnumerator(G, group):
         yield v
 
   for v in combine([group.zero()] * m, 0):
-    yield v
+    yield tuple(v)
 
 
-def testGroupConnectivityNaive(G_, group, forb_spanning_tree = False, list_all = False):
+def testGroupConnectivityNaive(G_, group, forb_spanning_tree = False,
+                               list_all = False, counts = False):
   E = [ (u, v, i) for i, (u, v, _) in enumerate(G_.edges()) ]
   G = DiGraph([G_.vertices(), E], format='vertices_and_edges')
   m = G.num_edges()
-  info = { 'graph': G, 'edges': E }
+  C = {}
+  info = { 'graph': G, 'edges': E, 'counts': C }
+
+  if counts:
+    list_all = True
 
   if forb_spanning_tree:
     T = G.min_spanning_tree()
@@ -65,28 +70,34 @@ def testGroupConnectivityNaive(G_, group, forb_spanning_tree = False, list_all =
   info['compatible_flows'] = flows
 
   Elems = [ x for x in group ]
-  forb = [ group.zero() ] * m
+  forb_ = [ group.zero() ] * m
 
   def forb_iter(e):
     if e >= len(T):
-      yield True
+      yield tuple(forb_)
       return
 
     for v in Elems:
-      forb[T[e]] = v
+      forb_[T[e]] = v
       for r in forb_iter(e + 1):
         yield r
 
-  def find_flow():
+  if not counts:
+    def find_flow(forb):
       for f in flows:
         if is_compatible(f, forb):
           return True
       return False
+  else:
+    def find_flow(forb):
+      c = len(filter(lambda f: is_compatible(f, forb), flows))
+      C[forb] = c
+      return (c > 0)
 
   ret = []
 
-  for _ in forb_iter(0):
-    if not find_flow():
+  for forb in forb_iter(0):
+    if not find_flow(forb):
       if not list_all:
         return (False, info, [forb])
       ret.append(list(forb))