I was wondering what the best way is (in Python) to iterate over partitions of a list of a given size.
Say, for example, we have the list [1,2,3,4,5] and we want k=3 partitions. A poor way of doing this would be to write:
lst = [1,2,3,4,5]
for i in range(1,len(lst)):
for j in range(i+1, len(lst)):
print lst[:i], lst[i:j], lst[j:]
This gives
[1], [2], [3,4,5]
[1], [2,3], [4,5]
...
[1,2,3], [4], [5]
But if I later wanted to iterate over k=4 partitions, then I would have to add a level of for loop nesting, which can't be done at runtime. Ideally, I'd like to write something like:
for part in partitions([1,2,3,4,5], k):
print part
Does anyone know the best way of doing this?
I would use the same idea as yours without pairwise:
from itertools import combinations
def partitions(items, k):
def split(indices):
i=0
for j in indices:
yield items[i:j]
i = j
yield items[i:]
for indices in combinations(range(1, len(items)), k-1):
yield list(split(indices))
I accomplished what I was trying to do by writing:
from itertools import tee, izip, combinations
def partitions(items, k):
N = len(items)
def pairwise(iterable): # Taken from itertools recipies
a, b = tee(iterable)
next(b, None)
return izip(a, b)
def applyPart(part, items):
lists = []
for l,h in pairwise([0] + part + [N]):
lists.append(items[l:h])
return lists
for part in combinations(range(1, N), k - 1):
yield applyPart(list(part), items)
This could be somewhat inefficient for larger lists, but it works:
from itertools import product, islice
def partitions(seq, k):
for c in product(xrange(1, len(seq)+1), repeat=k):
if sum(c) == len(seq):
it = iter(seq)
yield [list(islice(it, x)) for x in c]
for part in partitions([1,2,3,4,5], 3):
print part
Output:
[[1], [2], [3, 4, 5]]
[[1], [2, 3], [4, 5]]
[[1], [2, 3, 4], [5]]
[[1, 2], [3], [4, 5]]
[[1, 2], [3, 4], [5]]
[[1, 2, 3], [4], [5]]
For bigger lists what you need to is find all k sized subsets of range(1, len(sequence)+1) that sum to the length of the sequence and then slice the sequence based on them.
Related: http://www.algorithmist.com/index.php/Coin_Change
