包括选择元素的组合（Python）

Question

In order to make the set of all combinations of numbers 0 to x, with length y, we do:

list_of_combinations=list(combinations(range(0,x+1),y))
list_of_combinations=map(list,list_of_combinations)
print list_of_combinations

This will output the result as a list of lists.

For example, x=4, y=3:

[[0, 1, 2], [0, 1, 3], [0, 1, 4], [0, 2, 3], [0, 2, 4], [0, 3, 4], [1, 2, 3], [1, 2, 4], 
[1, 3, 4], [2, 3, 4]]

I am trying to do the above, but only outputting lists that have 2 members chosen beforehand.

For instance, I would like to only output the set of the combos that has 1 and 4 inside it. The output would then be (for x=4, y=3):

[[0, 1, 4], [1, 2, 4], [1, 3, 4]]

The best approach I have now is to make a list that is y-2 length with all numbers of the set without the chosen numbers, and then append the chosen numbers, but this seems very inefficient. Any help appreciated.

*Edit: I am doing this for large x and y, so I can't just write out all the combos and then search for the selected elements, I need to find a better method.

Answer 1

combinations() returns an iterable , so loop over that while producing the list:

[list(combo) for combo in combinations(range(x + 1), y) if 1 in combo]

This produces one list, the list of all combinations that match the criteria.

Demo:

>>> from itertools import combinations
>>> x, y = 4, 3
>>> [list(combo) for combo in combinations(range(x + 1), y) if 1 in combo]
[[0, 1, 2], [0, 1, 3], [0, 1, 4], [1, 2, 3], [1, 2, 4], [1, 3, 4]]

The alternative would be to produce y - 1 combinations of range(x + 1) with 1 removed, then adding 1 back in (using bisect.insort() to avoid having to sort afterwards):

import bisect

def combinations_with_guaranteed(x, y, *guaranteed):
    values = set(range(x + 1))
    values.difference_update(guaranteed)
    for combo in combinations(sorted(values), y - len(guaranteed)):
        combo = list(combo)
        for value in guaranteed:
            bisect.insort(combo, value)
        yield combo

then loop over that generator:

>>> list(combinations_with_guaranteed(4, 3, 1))
[[0, 1, 2], [0, 1, 3], [0, 1, 4], [1, 2, 3], [1, 2, 4], [1, 3, 4]]
>>> list(combinations_with_guaranteed(4, 3, 1, 2))
[[0, 1, 2], [1, 2, 3], [1, 2, 4]]

This won't produce as many combinations for filtering to discard again.

It may well be that for larger values of y and guaranteed numbers, just using yield sorted(combo + values) is going to beat repeated bisect.insort() calls.

Answer 2

This should do the trick:

filtered_list = filter(lambda x: 1 in x and 4 in x, list_of_combinations)

To make your code nicer (use more generators), I'd use this

combs = combinations(xrange(0, x+1), y)
filtered_list = map(list, filter(lambda x: 1 in x and 4 in x, combs))

If you don't need the filtered_list to be a list and it can be an iterable, you could even do

from itertools import ifilter, imap, combinations
combs = combinations(xrange(0, x+1), y)
filtered_list = imap(list, ifilter(lambda x: 1 in x and 4 in x, combs))
filtered_list.next()
>  [0, 1, 4]
filtered_list.next()
>  [1, 2, 4]
filtered_list.next()
>  [1, 3, 4]
filtered_list.next()
>  Traceback (most recent call last):
>    File "<stdin>", line 1, in <module>
>  StopIteration