CODEWITHSUNDEEP
pythonlist
Krishal
Krishal

Reputation: 43

Generate lists with elements less or equal to the given list

What is the most pythonic way to solve the following problem?

Given a list A, find all lists B, such that for i in range(len(A)): B[i] <= A[i]. Example of what I expect: 

#Input
A = [1,2,0]
#Output
B = [[0,0,0],
     [1,0,0],
     [1,1,0],
     [1,2,0],
     [0,1,0],
     [0,2,0]]

Thanks in advance!

Upvotes: 2

Views: 137

Answers (4)

levraininjaneer
levraininjaneer

Reputation: 1377

This also works:

A = [1,2,0]

answers = [[] for element in A]
answers[0] = [[number,] for number in range(A[0]+1)]

for i in range(len(A)-1):
    for possibility in answers[i]:
        for valid_number in range(A[i+1]+1):
            answers[i+1].append(possibility + [valid_number, ])

your_answer = answers[-1]

It could probably be written in a more compact way.

Upvotes: 0

Laurent H.
Laurent H.

Reputation: 6526

Here is a quite smart solution without using itertools but using recursive function call. It works whatever the size of the input list: 

A = [1,2,0]

def compute(A, B=[], n=0):
    for i in range(A[n]+1):
        A[n] = i
        if A not in B: B.append(A[:])
        if n+1 < len(A): compute(A, B, n+1)  # recursive call here
    return sorted(B)

B = compute(A)

print(B)
# [[0, 0, 0], [0, 1, 0], [0, 2, 0], [1, 0, 0], [1, 1, 0], [1, 2, 0]]

Upvotes: 1

RoadRunner
RoadRunner

Reputation: 26315

I guess you could first generate all unique combinations of these numbers:

>>> import itertools
>>> A = [1,2,0]
>>> combs = set(itertools.combinations(itertools.chain(*(range(col+1) for col in A)), len(A)))
>>> combs
{(0, 1, 1), (0, 2, 0), (1, 1, 0), (1, 2, 0), (0, 0, 2), (0, 1, 2), (1, 0, 0), (0, 0, 1), (1, 0, 1), (1, 1, 2), (0, 0, 0), (0, 1, 0), (1, 0, 2)}

Then pick the ones that follow your condition:

>>> sorted(comb for comb in combs if all(comb[i] <= x for i, x in enumerate(A)))
[(0, 0, 0), (0, 1, 0), (0, 2, 0), (1, 0, 0), (1, 1, 0), (1, 2, 0)]

But this is alot more inefficient than using itertools.product, since it has to generate all the combinations beforehand, and sort at the end to retain order.

Upvotes: 0

Sunitha
Sunitha

Reputation: 12005

You can use itertools.product to do this easily

>>> from itertools import product
>>> A
[1, 2, 0]
>>> B = list(product(*[list(range(e+1)) for e in A]))
>>> B
[(0, 0, 0), (0, 1, 0), (0, 2, 0), (1, 0, 0), (1, 1, 0), (1, 2, 0)]
>>>

If you want the o/p as list of list, convert the tuples to list

>>> B = [list(e) for e in B]
>>> B
[[0, 0, 0], [0, 1, 0], [0, 2, 0], [1, 0, 0], [1, 1, 0], [1, 2, 0]]
>>>

If you dont wan't to use itertools.product, you can have your custom implementation of product

>>> B = [[]]; 
>>> for t in [range(e+1) for e in A]:
...     B = [x+[y] for x in B for y in t]
... 
>>> B
[[0, 0, 0], [0, 1, 0], [0, 2, 0], [1, 0, 0], [1, 1, 0], [1, 2, 0]]
>>>

Upvotes: 7

Related Questions