kcomb

comb.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Write me a function that takes a collection S and an integer k, and returns a
collection of all collections such that each is a sub-collection of S with size k.
 - S is a set-like structure that guarantees these collections contain totally
    ordered and unique values; importantly, [1,2] is the same as [2,1].
 - You can make sets as needed.
 - They support indexing and slicing, as well as iteration.

>>> kcomb(S=[1, 2, 3, 4], k=2)
[ [1, 2], [1, 3], [1, 4],
  [2, 3], [2, 4],
  [3, 4] ]

"""


def kcomb(S, k):
    "Calculates the set of all subsets of S with size k."
    if k <= 0:
        return None
    elif k == 1:
        return [ [v] for v in S ]
    else:
        return [ [v]+suffix for i, v in enumerate(S) for suffix in kcomb(S[i+1:], k-1) ]