CODEWITHSUNDEEP
pythonlist-comprehension
Jithin Johnson
Jithin Johnson

Reputation: 372

Dynamic Nested Loops Python

This is the array I require:

N = 6
A = [[x,y,z] for x in range(N+1) for y in range(N+1) for z in range(N+1) if x+y+z== N]

Is there any other way of doing this by only specifying the variable N and 3 instead of x,y,z?
I tried []*3 but can't get the required output.

Upvotes: 1

Views: 225

Answers (2)

Timur Shtatland
Timur Shtatland

Reputation: 12347

Use itertools.combinations_with_replacement:

from itertools import combinations_with_replacement
n = 6
k = 3

# If you want a list of tuples:
lst = [item for item in list(combinations_with_replacement(range(n), k)) if sum(item) == n]
print(lst)
# [(0, 1, 5), (0, 2, 4), (0, 3, 3), (1, 1, 4), (1, 2, 3), (2, 2, 2)]

# If you want a list of lists:
lst = [list(item) for item in list(combinations_with_replacement(range(n), k)) if sum(item) == n]
print(lst)
# [[0, 1, 5], [0, 2, 4], [0, 3, 3], [1, 1, 4], [1, 2, 3], [2, 2, 2]]

Upvotes: 2

Tomerikoo
Tomerikoo

Reputation: 19404

You are looking for iterttols.product:

from itertools import product

N = 6
A = [(x,y,z) for x in range(N+1) for y in range(N+1) for z in range(N+1) if x+y+z== N]
B = [tup for tup in product(range(N+1), repeat=3) if sum(tup) == N]
print(A, B, A == B, sep='\n')

Gives:

[(0, 0, 6), (0, 1, 5), (0, 2, 4), (0, 3, 3), (0, 4, 2), (0, 5, 1), (0, 6, 0), (1, 0, 5), (1, 1, 4), (1, 2, 3), (1, 3, 2), (1, 4, 1), (1, 5, 0), (2, 0, 4), (2, 1, 3), (2, 2, 2), (2, 3, 1), (2, 4, 0), (3, 0, 3), (3, 1, 2), (3, 2, 1), (3, 3, 0), (4, 0, 2), (4, 1, 1), (4, 2, 0), (5, 0, 1), (5, 1, 0), (6, 0, 0)]
[(0, 0, 6), (0, 1, 5), (0, 2, 4), (0, 3, 3), (0, 4, 2), (0, 5, 1), (0, 6, 0), (1, 0, 5), (1, 1, 4), (1, 2, 3), (1, 3, 2), (1, 4, 1), (1, 5, 0), (2, 0, 4), (2, 1, 3), (2, 2, 2), (2, 3, 1), (2, 4, 0), (3, 0, 3), (3, 1, 2), (3, 2, 1), (3, 3, 0), (4, 0, 2), (4, 1, 1), (4, 2, 0), (5, 0, 1), (5, 1, 0), (6, 0, 0)]
True

Upvotes: 2

Related Questions