Search code examples
pythonlistpython-itertools

Unpacking a list of list of list


The current input I use in my program is a list of list of coordinates:

data = [["D_02", "10", "10", "20"],
         ["X_03", "10", "10", "20"]]

For ease of use and convenience, I would like to be able to pass a list of list of list as well:

have:

data = [[["B_01"], ["10"], ["10", "20"], ["10", "20"]],
             ["D_02", "10", "10", "20"],
             ["X_03", "10", "10", "20"]]

and then be able to "explode" it to make it like the following:

want:

data = [["B_01", "10", "10", "10"],
        ["B_01", "10", "10", "20"],
        ["B_01", "10", "20", "10"],
        ["B_01", "10", "20", "20"],
        ["D_02", "10", "10", "20"],
        ["X_03", "10", "10", "20"]]

I know it is possible to do so for a list of list with the product function of the itertools package:

import itertools
list(itertools.product(*data))

However,

  1. How can I do the same for a list of list of list ?
  2. How can I loop on the input to only explode the list of list of list ?

Solution

  • Use itertools.product only if the first item of row is a list:

    final = []
    for l in data:
        if isinstance(l[0], list):  # <- ["B01"]
            for c in itertools.product(*l[1:]):
                final.append(l[0] + list(c))
        else:
            final.append(l)
    
    >>> final
    [['B_01', '10', '10', '10'],
     ['B_01', '10', '10', '20'],
     ['B_01', '10', '20', '10'],
     ['B_01', '10', '20', '20'],
     ['D_02', '10', '10', '20'],
     ['X_03', '10', '10', '20']]