combinations.py

   1 """Assorted useful combinatorics.
   2 """
   3
   4 from __future__ import generators
   5
   6
   7 def xcombinations(items, n):
   8     """Iterate through ordered, length `n` subsets of `items`.
   9
  10     >>> for c in xcombinations([1, 2, 3], 2):
  11     ...     print c
  12     [1, 2]
  13     [1, 3]
  14     [2, 1]
  15     [2, 3]
  16     [3, 1]
  17     [3, 2]
  18     """
  19     if n==0: yield []
  20     else:
  21         for i in xrange(len(items)):
  22             for cc in xcombinations(items[:i]+items[i+1:],n-1):
  23                 yield [items[i]]+cc
  24
  25 def xunique_combinations(items, n):
  26     """Iterate through unordered, length `n` subsets of `items`.
  27
  28     >>> for c in xunique_combinations([1, 2, 3], 2):
  29     ...     print c
  30     [1, 2]
  31     [1, 3]
  32     [2, 3]
  33     """
  34     if n==0: yield []
  35     else:
  36         for i in xrange(len(items)):
  37             for cc in xunique_combinations(items[i+1:],n-1):
  38                 yield [items[i]]+cc
  39
  40 def xselections(items, n):
  41     """Iterate through all length `n` multisets of `items`.
  42
  43     >>> for c in xselections([1, 2, 3], 2):
  44     ...     print c
  45     [1, 1]
  46     [1, 2]
  47     [1, 3]
  48     [2, 1]
  49     [2, 2]
  50     [2, 3]
  51     [3, 1]
  52     [3, 2]
  53     [3, 3]
  54     """
  55     if n==0: yield []
  56     else:
  57         for i in xrange(len(items)):
  58             for ss in xselections(items, n-1):
  59                 yield [items[i]]+ss
  60
  61 def xpermutations(items):
  62     """Iterate through all permutations of `items`.
  63
  64     >>> for c in xpermutations([1, 2, 3]):
  65     ...     print c
  66     [1, 2, 3]
  67     [1, 3, 2]
  68     [2, 1, 3]
  69     [2, 3, 1]
  70     [3, 1, 2]
  71     [3, 2, 1]
  72     """
  73     return xcombinations(items, len(items))
  74