Search code examples
pythonpython-3.xlistlist-comprehensionnested-loops

How to expand this nested list comprehension in Python


Consider the below function:

It takes a list of lists as an input and finds all combinations of elements from each of those lists.

def product(llist):
    result = [[]]

    for lst in llist:
        result = [x + [y] for x in result for y in lst]

    return result

For eg:

product([[1,2], [3], [4,5]])

would return:

[[1, 3, 4], [1, 3, 5], [2, 3, 4], [2, 3, 5]]

I'm trying to understand how this function works and hence trying to expand the list comprehension.

Here's a try:

def product2(llist):
    result = [[]]
    for lst in llist:
        for x in result:
            result = []
            for y in lst:
                result.append(x+[y])

    return result

this doesn't give me the correct result, it returns:

[[2, 3, 4], [2, 3, 5]]

and I understand this incorrect result as per the definition of product2. But I'm not able to expand the original product function to understand how it is working.

Can someone elaborate the nested list comprehension in product function?


Solution

  • The list comprehension is creating a new list for each element in lst, and creates further new lists by combining this small single element list with all the lists that are already in result.

    So, for example if result = [ [1, 2], [3, 4] ] and lst = [10, 11], new result would be [[1, 2, 10], [1, 2, 11], [3, 4, 10], [3, 4, 11]].

    And it's doing this for each lst in llist:

    def product2(llist):
        result = [[]]
        for lst in llist:
            new_result = [] # let's manually build out new result
            for existing in result:            
                for element in lst:
                    new_result.append(existing + [element])
            result = new_result # update result for next lst
        return result