-2

Is there an easy way to write a generator to generate all lists that are k elements long, where each element ranges from 1 to n?

So if k=3 and n=3, it would generate

[1,1,1]
[1,1,2]
[1,1,3]
[1,2,1]
[1,2,2]
[1,2,3]
[1,3,1]
...
[3,3,1]
[3,3,2]
[3,3,3]
user115566
  • 41
  • 2

1 Answers1

4

Using itertools.product():

import itertools
gen = itertools.product(range(1, k+1), repeat=n)

This will generate tuples, if you want lists instead you can use itertools.imap() or a generator expression, for example:

gen = (list(t) for t in itertools.product(range(1, k+1), repeat=n))
Andrew Clark
  • 202,379
  • 35
  • 273
  • 306