Numpy where 沿轴 0 的条件语句

Question

我有一个Zc包含 n 个元素的一维向量，这些元素是二维数组。我想找到等于的每个二维数组的索引np.ones(Zc[i].shape)。

a = np.zeros((5,5))

b = np.ones((5,5))*4

c = np.ones((5,5))

d = np.ones((5,5))*2

Zc = np.stack((a,b,c,d))

for i in range(len(Zc)):

    a = np.ones(Zc[i].shape)

    b = Zc[i]

    if np.array_equal(a,b):

        print(i)

    else:

        pass 

返回2. 上面的代码有效并返回正确的答案，但我想知道是否有一种矢量化的方式来实现相同的结果？

Answer 1

离开 hpaulj 的评论：>>> allones = (Zc == np.array(np.ones(Zc[i].shape))).all(axis=(1,2))>>> np.where(allones)[0][0]2