沿轴0的条件语句的块状

时间:2018-09-12 19:51:47

标签: python numpy vector conditional

我有一个一维向量Zc,其中包含n个元素,它们是2D数组。我想找到每个等于np.ones(Zc[i].shape)的2D数组的索引。

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。上面的代码可以正常工作并返回正确的答案,但是我想知道是否有矢量化的方法来实现相同的结果?

1 个答案:

答案 0 :(得分:0)

不谈hpaulj的评论:

>>> allones = (Zc == np.array(np.ones(Zc[i].shape))).all(axis=(1,2))
>>> np.where(allones)[0][0]
2