Search code examples
pythonlistnested-listsdynamically-generated

Generating a list of dynamic length consisting of lists of dynamic lengths


I would like an efficient, neat way to generate the following list of sublists in Python (3.9), and it has me stumped. Child's play for most people I'm sure, but I'm obviously only at infant or toddler play.

Variables would be (with defaults used in brackets):

  • a is max number of entries in a given sublist (3)
  • x is starting value (25)
  • y is increment value (20)
  • z is maximum value (70)

Eventually I would add a check that the sublist is only appended to list if sum(sublist) is between two variable values, so I end up with a list of permutations that satisfy some criteria. Those criteria might include changing the value of y to keep the sheer number of iterations under control as well.

[
[25],
[45],
[65],
[25,25],
[45,25],
[65,25],
[25,45],
...
[45,65],
[65,65],
[25,25,25],
[45,25,25],
...
[45,65,65],
[65,65,65]
]

Solution

  • You can use itertools.product() for each subset of length, and then itertools.chain() to concatenate the lists together.

    def get_lists(a, x, y, z, m=float('-inf'), n=float('inf')):
        v = list(range(x, z, y))
        return list(itertools.chain(*(
            (
                list(reversed(p))
                for p in itertools.product(*([v] * i)) 
                if m < sum(p) < n
            ) for i in range(1, a + 1)
        )))
    

    This also accounts for the lower and upper bounds if you want, with optional arguments m and n. If not given, well, everything should fit between infinity and negative infinity.

    A couple of tricks are used here.

    • First of all, itertools.product() will display in the opposite order that you want (it varies the last element first, whereas you vary the first element first). So we feed it into reversed(). itertools.product() also normally yields tuples instead of lists, but since we have to cast the iterator reversed() returns to a list anyway, it's no trouble. If you don't care about this, you can omit the innermost list(reversed(p)), and just leave it at p. Or if you don't care about filtering the input after all, you could just remove the inner comprehension entirely.
    • Second, in both itertools.product() and itertools.chain(), we use the argument-unpacking operator * to pass a whole list as positional arguments. For itertools.product(), this is just duplicating the 'eligible values' range as many times as necessary, whereas for itertools.chain() it's stitching together the individual outputs of each itertools.product(). This idiom is required for a lot of itertools abuse.

    Otherwise this is just a few nested comprehensions. Technically you could do it in a one-liner, but I find that putting v in a variable and doing [v] * i is nicer than doing (range(x, z, y) for _ in range(i)).