Numpy где условный оператор по оси 0 - PullRequest
0 голосов
/ 12 сентября 2018

У меня есть 1D вектор 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.Приведенный выше код работает и возвращает правильный ответ, но я хочу знать, есть ли векторизованный способ достижения того же результата?

1 Ответ

0 голосов
/ 12 сентября 2018

Исходя из комментария hpaulj:

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