如何从列表/ ndarray获取索引？

Question

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.

Answer 1

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

1 2
2 1
2 2
3 0
3 1
3 2

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.

Answer 2

You could just compare the array to your value and use where.

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

(array([1, 2, 2, 3, 3, 3], dtype=int64), array([2, 1, 2, 0, 1, 2], dtype=int64))

To get your tuples

list(zip(*np.where(a>2)))

[(1, 2), (2, 1), (2, 2), (3, 0), (3, 1), (3, 2)]

Answer 3

I have made it out, the code should be:

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 zip(*(np.nonzero(a>2))):
    if  r == cur_row:
        col_list.append(c)
    else:
        cur_row = r
        shuffle(col_list)
        if len(col_list) == 0:
            out[r-1] = -1
        else:
            out[r-1] = col_list[0]
        col_list = []
        col_list.append(c)

shuffle(col_list)
if len(col_list) == 0:
    out[len(out)-1] = -1
else:
    out[len(out)-1] = col_list[0]

The part in the end but outside the forloop is to make sure that the last row will be taken care of.

It works in my case.