How to get indices from a list/ndarray?

I have a list which looks like:

[[0,1,2], [1,2,3], [2,3,4], [3,4,5]]

I can make it to an array like:

array([[0,1,2],
       [1,2,3],
       [2,3,4],
       [3,4,5]])

So all together I have 4 rows and each row has 3 columns. Now I want to find the indices of all the elements which are greater than 2, so for the whole matrix, the indices should be:

((1,2),(2,1),(2,2),(3,1),(3,2),(3,3))

Then for each row, I will randomly picked out a col index which indicates a value greater than 2. Now my code is like:

a = np.array([[0,1,2],[1,2,3],[2,3,4],[3,4,5]]
out = np.ones(4)*-1
cur_row = 0
col_list = []
for r,c in np.nonzero(a>2):
    if r == cur_row:
        col_list.append(c)
    else:
        cur_row = r
        shuffled_list = shuffle(col_list)
        out[r-1] = shuffled_list[0]
        col_list = []
        col_list.append(c)

I hope to get a out which looks like:

array([-1, 2, 1, 2])

However, now when I run my code, it shows

ValueError: too many values to unpack

Anyone knows how I fix this problem? Or how should I do to achieve my goal? I just want to run the code as fast as possible, so any other good ideas is also more than welcome.

Solution

Try this.

import numpy as np

arr = np.array([[0,1,2],
               [1,2,3],
               [2,3,4],
               [3,4,5]])
indices = np.where(arr>2)

for r, c in zip(*indices):
    print(r, c)

Prints

So, it should work. You can use itertools.izip as well, it would even be a better choice in this case.

A pure numpy solution (thanks to @AshwiniChaudhary for the proposition):

for r, c in np.vstack(np.where(arr>2)).T:
    ...

though I'm not sure this will be faster than using izip or zip.