git.notmuchmail.org Git - notmuch/blob - test/gen-threads.py

   1 # Generate all possible single-root message thread structures of size
   2 # argv[1].  Each output line is a thread structure, where the n'th
   3 # field is either a number giving the parent of message n or "None"
   4 # for the root.
   5 import sys
   6 from itertools import chain, combinations
   7
   8 def subsets(s):
   9     return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
  10
  11 nodes = set(range(int(sys.argv[1])))
  12
  13 # Queue of (tree, free, to_expand) where tree is a {node: parent}
  14 # dictionary, free is a set of unattached nodes, and to_expand is
  15 # itself a queue of nodes in the tree that need to be expanded.
  16 # The queue starts with all single-node trees.
  17 queue = [({root: None}, nodes - {root}, (root,)) for root in nodes]
  18
  19 # Process queue
  20 while queue:
  21     tree, free, to_expand = queue.pop()
  22
  23     if len(to_expand) == 0:
  24         # Only print full-sized trees
  25         if len(free) == 0:
  26             print(" ".join(map(str, [msg[1] for msg in sorted(tree.items())])))
  27     else:
  28         # Expand node to_expand[0] with each possible set of children
  29         for children in subsets(free):
  30             ntree = {child: to_expand[0] for child in children}
  31             ntree.update(tree)
  32             nfree = free.difference(children)
  33             queue.append((ntree, nfree, to_expand[1:] + tuple(children)))