Numpy 使用 2D 数组作为 3D 数组的类似高度图的索引

Question

我想使用 2D 数组作为 3D 数组的索引作为高度图来索引 3D 数组的轴 0。这样做有效率 "numpy-way" 吗？在我的例子中，我想在每个相应的支柱两个零中将所有内容设置为等于或大于高度图的高度。例子： 3D 阵列：

[[[1, 1, 1],
  [1, 1, 1],
  [1, 1, 1]],
 [[1, 1, 1],
  [1, 1, 1],
  [1, 1, 1]],
 [[1, 1, 1],
  [1, 1, 1],
  [1, 1, 1]],
 [[1, 1, 1],
  [1, 1, 1],
  [1, 1, 1]]]

二维数组（高度图）：

[[0, 1, 2],
 [2, 3, 4],
 [2, 0, 0]]

期望的输出：

[[[0, 1, 1],
  [1, 1, 1],
  [1, 0, 0]],
 [[0, 0, 1],
  [1, 1, 1],
  [1, 0, 0]],
 [[0, 0, 0],
  [0, 1, 1],
  [0, 0, 0]],
 [[0, 0, 0],
  [0, 0, 1],
  [0, 0, 0]]]

到目前为止，我已经使用 for python 循环实现了这一点，如

for y in range(arr2d.shape[0]):
    for x in range(arr2d.shape[1]):
        height = arr2d[y, x]
        arr3d[height:, y, x] = 0

但这似乎非常低效，我觉得可能有更好的方法来做到这一点。

Answer 1

从快速填充数组的方法中汲取灵感：

In [104]: (np.arange(4)[:,None,None]<arr2d).astype(int)                                          
Out[104]: 
array([[[0, 1, 1],
        [1, 1, 1],
        [1, 0, 0]],

       [[0, 0, 1],
        [1, 1, 1],
        [1, 0, 0]],

       [[0, 0, 0],
        [0, 1, 1],
        [0, 0, 0]],

       [[0, 0, 0],
        [0, 0, 1],
        [0, 0, 0]]])

Numpy 使用 2D 数组作为 3D 数组的类似高度图的索引

Numpy Use 2D array as heightmap-like index for 3D array

python

arrays

indexing

numpy

heightmap