Numpy where 沿轴 0 的条件语句

Question

我有一个Zc包含 n 个元素的一维向量，这些元素是二维数组。我想找到等于的每个二维数组的索引np.ones(Zc[i].shape)。a = np.zeros((5,5))b = np.ones((5,5))*4c = np.ones((5,5))d = np.ones((5,5))*2Zc = 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