Python Code optimizer for list of list choices conbination

Viewed 50

I have python code to find all the possible choices from the list of lists. But when I have a big size list it gives a larger possible dataset. My function is,

def yield_combos(lst):
    if not lst:
        yield []
    else:
        for el in lst[0]:
            for combo in yield_combos(lst[1:]):
                yield [el] + combo

Above code to the input of,

l=[[1,2,3],[4,5,6],[0,2,7]]

output is,

[1, 4, 0]
[1, 4, 2]
[1, 4, 7]
[1, 5, 0]
[1, 5, 2]
[1, 5, 7]
[1, 6, 0]
[1, 6, 2]
[1, 6, 7]
[2, 4, 0]
[2, 4, 2]
[2, 4, 7]
[2, 5, 0]
[2, 5, 2]
[2, 5, 7]
[2, 6, 0]
[2, 6, 2]
[2, 6, 7]
[3, 4, 0]
[3, 4, 2]
[3, 4, 7]
[3, 5, 0]
[3, 5, 2]
[3, 5, 7]
[3, 6, 0]
[3, 6, 2]
[3, 6, 7]

Now I want to reduce the output list with a condition like, all the possible list should include at least one list first item. Then it should remove below item from the list,

[2,5,2]
[2,5,7]
[2,6,2]
[2,6,7]
[3,5,2]
[3,5,7]
[3,6,2]
[3,6,7]

This will help to optimize my algorithm for a large list. How can I do this?

2 Answers

You can modify the function as such:

def yield_combos(lst, includes_head_item=False):
    if not lst:
        if not includes_head_item:
            return
        yield []
    else:
        for el_index, el in enumerate(lst[0]):
            for combo in yield_combos(lst[1:], includes_head_item or el_index == 0):
                yield [el] + combo

We use the variable includes_head_item to keep track of whether any sublist's head item has been included in the combo. If we reach the end with that condition still false, the code is structured such that we don't yield, removing that combination from the generated ones.

If you wanted to make such a function for other indices as well (for example, ensuring that each combination includes atleast one element at the end of a sublist), you could do this:

def yield_combos(lst, included_index=0, includes_item=False):
    if not lst:
        if not includes_item:
            return
        yield []
    else:
        for el_index, el in enumerate(lst[0]):
            for combo in yield_combos(lst[1:], included_index, includes_item or el_index == included_index):
                yield [el] + combo

To accomplish the above example for including atleast one last element, you would call it as such:

yield_combos(..., 2)

the combination list can be achieved with itertools.product

For example product('ABCD', 'xy') gives Ax Ay Bx By Cx Cy Dx Dy, thus for your example:

itertools.product([1,2,3],[4,5,6],[0,2,7])

As for the final conditioning on zeros, it is simply excluding 0 from your initial list! (isn't it?!) but you can also do the conditioning with later if or assert checks

Related