CODEWITHSUNDEEP
python
Suraj Kothari
Suraj Kothari

Reputation: 1632

Better method to iterate over 3 lists

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.

Upvotes: 7

Views: 546

Answers (3)

Olivier Melançon
Olivier Melançon

Reputation: 22324

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)

Upvotes: 11

CezarySzulc
CezarySzulc

Reputation: 2009

You can use iterate over matrix. Its ndenumerate function from numpy packet. If width and height are range of int, you can skip creating list. Just define the size of this lists. 

width = 10
height = 5
k = [18,20,11]

kIndex = 0

for (i,j), value in np.ndenumerate(np.ones((height,width))):
    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)

Upvotes: 0

Christian Sloper
Christian Sloper

Reputation: 7510

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))

Upvotes: 1

Related Questions