Numpy Использование 2D-массива в качестве индекса, подобного карте высот, для 3D-массива

Question

Я хочу использовать 2D-массив в качестве индекса для 3D-массива в качестве карты высот для индексации оси 0 3D-массива. Есть ли эффективный "numpy -way" для этого? В моем примере я хочу установить все на равную или большую высоту карты высот в каждом соответствующем столбе два нуля. Пример: 3D Array:

[[[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]]]

2D Array (heightmap):

[[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]]]

До сих пор я реализовал это с помощью для python l oop как в

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]]])