这个问题与这个提问非常相似如何在TensorFlow中找到第一个匹配元素的索引

我尝试了那个问题的解决方案，但不同的是，val不是一个单一的数字，而是一个张量

例如

np.array([1, 1, 1],
         [1, 0, 1],
         [0, 0, 1])
val = np.array([1, 0, 1])
这里会发生一些TensorFlow的魔法！
result = 1

我知道我可以使用while循环，但这看起来很乱。我可以尝试使用映射函数，但有没有更优雅的方法呢？

回答：

这里是一种方法 –

(arr == val).all(axis=-1).argmax()

示例运行 –

In [977]: arr
Out[977]: array([[1, 1, 1],
       [1, 0, 1],
       [0, 0, 1]])
In [978]: val
Out[978]: array([1, 0, 1])
In [979]: (arr == val).all(axis=1).argmax()
Out[979]: 1

使用views可能会更高效 –

# https://stackoverflow.com/a/44999009/ @Divakar
def view1D(a): # a是数组
    a = np.ascontiguousarray(a)
    void_dt = np.dtype((np.void, a.dtype.itemsize * a.shape[1]))
    return a.view(void_dt).ravel()
out = (view1D(arr) == view1D(val[None])).argmax()

扩展到n维情况

扩展到n维数组的情况需要几个额外的步骤 –

def first_match_index_along_axis(arr, val, axis):
        s = [None]*arr.ndim
    s[axis] = Ellipsis
    mask = val[np.s_[s]] == arr
    idx = mask.all(axis=axis,keepdims=True).argmax()
    shp = list(arr.shape)
    del shp[axis]
    return np.unravel_index(idx, shp)

示例运行 –

In [74]: arr = np.random.randint(0,9,(4,5,6,7))
In [75]: first_match_index_along_axis(arr, arr[2,:,1,0], axis=1)
Out[75]: (2, 1, 0)
In [76]: first_match_index_along_axis(arr, arr[2,1,3,:], axis=3)
Out[76]: (2, 1, 3)