Tuesday, August 13, 2013

Python functions for enumerating combinatorics

Here are some useful enumerating functions in Python which have to do with generating permutations and combinations.

Enumerating all permutations of a sequence

This is a function which will give you all the different ways an ordered sequence can be ordered. It takes each element from the sequence and then recursively finds all the permutations of the sequence without that element. For every sub permutation the function will then prefix it with the removed element. The base case is that a sequence of just one element can only be ordered as itself.
def permutations(initial_permutation):
 if len(initial_permutation) == 1:
  yield initial_permutation
  for i in range(len(initial_permutation)):
   ith_element = initial_permutation[i:i+1]
   without_ith_element = initial_permutation[:i]+initial_permutation[i+1:]
   for sub_perm in permutations(without_ith_element):
    yield ith_element + sub_perm

['ABCD', 'ABDC', 'ACBD', 'ACDB', 'ADBC', 'ADCB', 'BACD', 'BADC', 'BCAD', 'BCDA', 'BDAC', 'BDCA', 'CABD', 'CADB', 'CBAD', 'CBDA', 'CDAB', 'CDBA', 'DABC', 'DACB', 'DBAC', 'DBCA', 'DCAB', 'DCBA']

Enumerating all unordered pairs in a sequence

This is a function which will give you all the different ways an unordered pair of elements can be taken from a sequence. It is a simple case of the next function. It simply runs two nested for loops, each of which selects an index of an element in such as way that the two selected indices are unique.
def pairs(seq):
 for i in range(0, len(seq)-1):
  for j in range(i+1, len(seq)):
   yield (seq[i], seq[j])

[('A', 'B'), ('A', 'C'), ('A', 'D'), ('B', 'C'), ('B', 'D'), ('C', 'D')]

Enumerating all unordered subsets of a given length in a sequence

This is the general case of the previous function, where instead of pairs it will enumerate tuples of a given length. It simply replaces the nested for loops with recursive for loops, each time the number of elements to choose goes down by one until the function has to choose 0 elements at which point it returns an empty tuple.
def choices(seq, choice_len):
 if choice_len == 0:
  yield ()
  for i in range(0, len(seq)-(choice_len-1)):
   for sub_choice in choices(seq[i+1:], choice_len-1):
    yield (seq[i],) + sub_choice
list(choices("ABCD", 3))

[('A', 'B', 'C'), ('A', 'B', 'D'), ('A', 'C', 'D'), ('B', 'C', 'D')]