Search code examples
pythonpython-itertools

Is there a way to get the combinations of a list of just one element to all the other elements of a list?


I am trying to generate the combinations of the first element of a list to all the other elements. Doing this with a small list is not an issue by first obtaining all the possible combinations and then discarding the ones I don't need. Example:

import numpy as np
import itertools

lst = ['a','b','c','d','e','f']
element_of_interest = 'a'

all_combinations = np.array(list(itertools.combinations(lst, 3)))

idx = np.where(all_combinations == element_of_interest)
combinations_of_interest = all_combinations[idx[0]]

print(combinations_of_interest)

>>> [['a' 'b' 'c']
 ['a' 'b' 'd']
 ['a' 'b' 'e']
 ['a' 'b' 'f']
 ['a' 'c' 'd']
 ['a' 'c' 'e']
 ['a' 'c' 'f']
 ['a' 'd' 'e']
 ['a' 'd' 'f']
 ['a' 'e' 'f']]

But in my application (list of 290 elements, r = 4) this way of doing it is not good enough since it's very inefficient.

Can anyone think of a better way of doing it?


Solution

  • Take the element out of the list and then add it later to all the rest combinations:

    In [9]: import itertools
       ...:
       ...: lst = ['a','b','c','d','e','f']
       ...: element_of_interest = 'a'
    
    In [11]: out = [['a'] + list(i) for i in itertools.combinations(lst[1:], 2)]
    
    In [12]: out
    Out[12]:
    [['a', 'b', 'c'],
     ['a', 'b', 'd'],
     ['a', 'b', 'e'],
     ['a', 'b', 'f'],
     ['a', 'c', 'd'],
     ['a', 'c', 'e'],
     ['a', 'c', 'f'],
     ['a', 'd', 'e'],
     ['a', 'd', 'f'],
     ['a', 'e', 'f']]