通过2D阵列遮罩3D阵列以进行切片,而无需for循环

问题描述

我有这样的东西

import numpy as np 

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

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

有没有一种方法可以避免double for循环?

解决方法

创建具有外部比较范围的遮罩,然后分配-

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