沿轴0的Numpy where条件语句

Zc 包含n个二维数组元素。我想找到每个2D数组的索引 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 . 上面的代码可以工作并返回正确的答案,但是我想知道是否有一种矢量化的方法来实现相同的结果?