Маска 3D-массива с помощью 2D-массива для нарезки без цикла for

#python #numpy #vectorization #numpy-slicing

#python #numpy #векторизация #numpy-нарезка

Вопрос:

У меня есть что-то вроде этого

 import numpy as np 

array_3D =  np.random.rand(3,3,3) 
array_2D = np.random.randint(0, 3 , (3,3)) 

for i in range(3):
    for j in range(3): 
        array_3D[:, i, j][:array_2D[i, j]]=np.nan  
  

Есть ли способ сделать это без двойного цикла for?

Ответ №1:

Создайте маску с внешним сравнением по диапазону, а затем назначьте —

 mask = np.less.outer(np.arange(len(array_3D)), array_2D)  
array_3D[mask] = np.nan