Search code examples
pythonarraysnumpyindexingheightmap

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

I want to use a 2D array as an index for a 3D array as a heightmap to index axis 0 of the 3D array. Is there an efficient "numpy-way" of doing this? In my example I want to set everything at equal or greater height of the heightmap in each corresponding pillar two zero. Example: 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]]

Desired output: 

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

So far I have implemented this with a for python loop as in 

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

but this seems very ineffecient and I feel like there might be a way better way to do this.

Solution

  • Drawing inspiration from an fast way of padding arrays:

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