Как эффективно использовать индексный массив в качестве маски, чтобы превратить массив numpy в булевский массив?

Question

У меня есть Numpy массив так:

>>> I 
array([[ 1., 0., 2., 1., 0.], 
     [ 0., 2., 1., 0., 2.]]) 

и массив А как это:

>>> A = np.ones((2,5,3)) 

Я хотел бы получить следующую матрицу:

>>> result 
array([[[ False, False, True], 
     [ False, True, True], 
     [ False, False, False], 
     [ False, False, True], 
     [ False, True, True]], 

     [[ False, True, True], 
     [ False, False, False], 
     [ False, False, True], 
     [ False, True, True], 
     [ False, False, False]]], dtype=bool) 

Это лучше пояснить на примере:
I[0,0] = 1 -> result[0,0,:2] = False и result[1,1,2:] = True
I[1,0] = 0 -> result[1,1,0] = False и result[1,1,1:] = True

Вот моя текущая реализация (правильный):

result = np.empty((A.shape[0], A.shape[1], A.shape[2])) 
r = np.arange(A.shape[2]) 
for i in xrange(A.shape[0]): 
    result[i] = r > np.vstack(I[i]) 

print result.astype(np.bool) 

Есть ли способ, чтобы реализовать в более быстрый способ (избегая цикл)?

Спасибо!

Answer 1

Вам просто нужно добавить еще одно измерение к I, так что вы можете транслировать r правильно:

result = r > I.reshape(I.shape[0],I.shape[1],1) 

например

In [41]: r>I.reshape(2,5,1) 
Out[41]: 
array([[[False, False, True], 
     [False, True, True], 
     [False, False, False], 
     [False, False, True], 
     [False, True, True]], 

     [[False, True, True], 
     [False, False, False], 
     [False, False, True], 
     [False, True, True], 
     [False, False, False]]], dtype=bool)