Unpack `np.unravel_index()` in for loop

Question

I'm trying to directly unpack the tuple returned from np.unravel_index() in a for loop definition, but I'm running into the following issue:

for i,j in np.unravel_index(range(len(10)), (2, 5)):
    print(i,j)

returns:

ValueError: too many values to unpack (expected 2)

I can solve the issue by doing:

idx = np.unravel_index(range(len(10)), (2, 5))
for i, j in zip(idx[0], idx[1]):
    print(i, j)

but I really feel I should be able to do it all in the for loop assignment.

I have looked through StackOverflow and found nothing that could help me with my specific question.

Solution: As the accepted solution I think that this does exactly what I want i.e., unpack directly in the for loop assignment and without previous knowledge of the dimensions of idx:

for i, j in zip(*np.unravel_index(range(len(10)), (2, 5)):
    print(i, j)

`unravel_index` returns a tuple of arrays. You've done correctly to `zip` them. There is an idiom to "transpose" nested lists/tuples: [python matrix transpose and zip](https://stackoverflow.com/questions/10169919/python-matrix-transpose-and-zip) — mkrieger1, Nov 26 '21 at 10:07

score 1 · Accepted Answer · answered Nov 26 '21 at 17:54

Your idx is a tuple of arrays:

In [559]: idx = np.unravel_index(np.arange(5),(2,5))
In [560]: idx
Out[560]: (array([0, 0, 0, 0, 0]), array([0, 1, 2, 3, 4]))

The tuple works great for indexing, e.g. data[idx]. In fact that's its intended purpose.

Your zip turns that into a list/iteration on 2 element tuples:

In [561]: list(zip(*idx))
Out[561]: [(0, 0), (0, 1), (0, 2), (0, 3), (0, 4)]

np.transpose can also turn it into a (n,2) array:

In [562]: np.transpose(idx)
Out[562]: 
array([[0, 0],
       [0, 1],
       [0, 2],
       [0, 3],
       [0, 4]])

Iteration on [562] will be just as slow as the iteration on the zip, possibly slower. But if you don't need to iterate, [562] may be better.

Notice I used zip(*idx) above, so your expression could written as:

for i, j in zip(*np.unravel_index(range(len(neuron_sample)), (2, 5))):
    print(i, j)

Unpack `np.unravel_index()` in for loop

1 Answers1