numpy 正在获取每行中第一个匹配项的索引

xwmevbvl  于 2023-06-23  发布在  其他
关注(0)|答案(3)|浏览(92)

我有一个充满布尔值的数组:

array([[[ True,  True, False, False, False, False],
        [ True, False, False, False, False, False]],

       [[False, False, True, False, True, False],
        [ True, False, False, False, False, False]],

       [[ True, False, False, False, False, False],
        [ True, False, False, False, False, False]]], dtype=bool)

我想得到每一行每一列中第一次出现True的索引,所以答案会是这样的:

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

有没有一个简单而快速的方法来做到这一点?

czq61nw1

czq61nw11#

现在不能测试,但我认为这应该工作

arr.argmax(axis=1).T

argmax在numpy 1.9中的bools短路,因此对于此用例,它应该优于wherenonzero

EDIT好的,所以上面的解决方案不起作用,但是argmax的方法仍然有用:

In [23]: mult = np.product(arr.shape[:-1])

In [24]: np.column_stack(np.unravel_index(arr.shape[-1]*np.arange(mult) +
   ....:                                  arr.argmax(axis=-1).ravel(),
   ....:                                  arr.shape))
Out[24]:
array([[0, 0, 0],
       [0, 1, 0],
       [1, 0, 2],
       [1, 1, 0],
       [2, 0, 0],
       [2, 1, 0]])
3ks5zfa0

3ks5zfa02#

似乎您希望np.where()the solution of this answer结合使用来查找唯一的行:

b = np.array(np.where(a)).T
#array([[0, 0, 0],
#       [0, 0, 1],
#       [0, 1, 0],
#       [1, 0, 2],
#       [1, 0, 4],
#       [1, 1, 0],
#       [2, 0, 0],
#       [2, 1, 0]], dtype=int64)
c = b[:,:2]
d = np.ascontiguousarray(c).view(np.dtype((np.void, c.dtype.itemsize * c.shape[1])))
_, idx = np.unique(d, return_index=True)

b[idx]
#array([[0, 0, 0],
#       [0, 1, 0],
#       [1, 0, 2],
#       [1, 1, 0],
#       [2, 0, 0],
#       [2, 1, 0]], dtype=int64)
arknldoa

arknldoa3#

添加到其他答案上,如果你想让它给予你第一个为True的列的索引,并返回n,其中n = # cols in a,如果一行不包含True:

first_occs = np.argmax(a, axis=1)
all_zeros = ~a.any(axis=1).astype(int)
first_occs_modified = first_occs + all_zeros * a.shape[1]

相关问题