使用np.where查找2D数组中的元素索引会产生ValueError

Question

I'm trying to use np.where to find the index of an element in an array, specifically row number

I have an array of say size 1000 x 6, named 'table'. The first element in each row is a 2 x 2 string array, and the rest are 0s. Eg. a 5 x 6 example of elements in 'table':

    [['s',' ']   0 0 0 0 0
     [' ',' ']]
    [[' ',' ']   0 0 0 0 0
     [' ','a']]
    [[' ',' ']   0 0 0 0 0
     [' ',' ']]         
    [['p',' ']   0 0 0 0 0
     [' ',' ']]
    [[' ',' ']   0 0 0 0 0
     ['b',' ']]  

The 2x2 arrays are all different, and I want to get the index, in particular the row number, of the one containing a specific 2x2 in my large table.

Eg. say I have

    grid = [['s',' ']   
            [' ',' ']]

I would like my code to return [0][0]

I have tried this:

    i,j = np.where(table == grid)

and also

    i,j = np.where(np.all(table == grid))

and i get the following error:

    ValueError: not enough values to unpack (expected 2, got 1)

Using a single value eg.

    index = np.where(table == grid) 

does not result in an error, but print(index) will output an empty array:

    (array([], dtype=int64),)

From similar questions on Stack Overflow I can't seem to figure out how this error applies to mine and I've been staring at it for ages

Any help would be much appreciated

Answer 1

Setup:

b = np.array([['s','t'],['q','r']])
c = np.array([['s',' '],[' ',' ']])
a = np.array([[c,0,0,0,0,0],
              [c,0,0,0,0,0],
              [c,0,0,0,0,0],
              [c,0,0,0,0,0],
              [b,0,0,0,0,0],
              [c,0,0,0,0,0],
              [c,0,0,0,0,0],
              [c,0,0,0,0,0],
              [c,0,0,0,0,0]])

Assuming your only interested in column zero; write a function that will test each item in a one-d array. And apply it to column zero

def f(args):
    return [np.all(thing==b) for thing in args]

>>> np.apply_along_axis(f,0,a[:,0])
array([False, False, False, False,  True, False, False, False, False])
>>> 

Use np.where on the result

>>> np.where(np.apply_along_axis(f,0,a[:,0]))
(array([4], dtype=int64),)

Or following the note in the numpy.where docs:

>>> np.asarray(np.apply_along_axis(f,0,a[:,0])).nonzero()
(array([4], dtype=int64),)

---

As @hpaulj points out np.apply_along_axis is not necessary. so

>>> [np.all(thing == b) for thing in a[:,0]]
[False, False, False, False, True, False, False, False, False]

>>> np.asarray([np.all(thing == b) for thing in a[:,0]]).nonzero()
(array([4], dtype=int64),)

And without the Python iteration:

>>> (np.stack(a[:,0])==b).all(axis=(1,2))
array([False, False, False, False,  True, False, False, False, False])

>>> (np.stack(a[:,0])==b).all(axis=(1,2)).nonzero()
(array([4], dtype=int64),)

Answer 2

Here is a solution using vectorize

a = np.array( [np.array([['s',' '],[' ',' ']])  , 0, 0, 0, 0, 0 ])
grid = np.array([['s',' '],[' ',' ']]) 

vfunc = np.vectorize(lambda x: np.all(grid == x))
np.argwhere(vfunc(a))