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

[Anonymous]

Я хочу использовать 2D-массив в качестве индекса для 3D-массива в качестве карты высот для индексации оси 0 3D-массива. Есть ли эффективный способ сделать это с помощью NumPy? Я хочу установить все, что находится на равной или большей высоте карты высот в каждом соответствующем столбе, на ноль.
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]]]

2D-массив (карта высот):

Код: Выделить всё

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

Это кажется неэффективным, но может быть есть лучший способ.