Better method to iterate over 3 lists

Question

I am creating a program which iterates over the width and height of an image as well as makes use of a set of keys.

Here is an example:

width = [0,1,2,3,4,6,7,8,9]
height = [0,1,2,3,4]
keys = [18,20,11]

The width and height are a range of integers up to the size of the width and height. The keys are any set of numbers (actually ASCII values) but are not ordered numbers.

I would like the output to be like this:

0 0 18
0 1 20
0 2 11
0 3 18
0 4 20
1 0 11
1 1 18
. . ..
9 0 20
9 1 11
9 2 18
9 3 20
9 4 11

As you can see, the width and height can be produced with nested for loops, whereas the keys cycle between each other.

Here is my solution:

w = [0,1,2,3,4,6,7,8,9]
h = [0,1,2,3,4]
k = [18,20,11]

kIndex = 0

for i in w:
    for j in h:
        print(i,j,k[kIndex])
        # Cycle through the keys index.
        # The modulo is used to return to the beginning of the keys list
        kIndex = (kIndex + 1) % len(k)

Actually it works as intended, however, I would like a more efficient way to do the above instead of using an incremental variable for the index position of the keys list.

I don't mind the nested for loop, if that has to be used, but the index keys variable annoys me as it seems like the code won't work without it, but at the same time isn't really pythonic.

Answer 1

You can use cycle from itertools.

from itertools import cycle, product

width = [0,1,2,3,4,6,7,8,9]
height = [0,1,2,3,4]
keys = [18,20,11]


c = cycle(keys)

for w,h in product(width,height):
    print(w,h,next(c))

Answer 2

You can use itertools.product to get the product of your width and height, that is your whole grid. Then, you want to cycle over the keys, thus use itertools.cycle. You finally zip those together and get the desired result.

You can make this a generator using yield for memory efficieny.

from itertools import product, cycle

def get_grid(width, height, keys):
    for pos, key in zip(product(width, height), cycle(keys)):
        yield (*pos, key)

Or if you do not want a generator.

out = [(*pos, key) for pos, key in zip(product(width, height), cycle(keys))]

Example

width = [0,1,2,3,4,6,7,8,9]
height = [0,1,2,3,4]
keys = [18,20,11]

for triple in get_grid(width, height, keys):
    print(triple)

Output

(0, 0, 18)
(0, 1, 20)
(0, 2, 11)
(0, 3, 18)
(0, 4, 20)
(1, 0, 11)
(1, 1, 18)
...

As a sidenote, notice that you could replace the lists defining width and height by ranges.

width = range(10)
height = range(5)