Is there a way to manipulate a 3D array with a 1D using vectorization?

In code below I've presented three ways to achieve the required outcome with their runtime comparisons:

1. Your for loop.
2. NumPy's apply_along_axis function
3. Numba's @jit decorator, which basically compiles your for loop.

import numpy as np 
from numba import jit

def set_ones(arr_tup):
    arr_tup[0][0][:arr_tup[1]] = 1

@jit
def compiled_func(data, num):
    for i in range(len(num)): 
        data[:, 0][i, 0:num[i]] = 1

data = np.zeros((300, 4, 5))
num = np.random.randint(2, 5, 300)

data_num = np.array(list(zip(data, num)))

print(f'\n> for loop:')
%timeit for i in range(len(num)): data[:, 0][i, 0:num[i]] = 1

print(f'\n> apply_along_axis Loop:')
%timeit np.apply_along_axis(set_ones, 1, data_num)

print(f'\n> Compiled Loop:')
%timeit compiled_func(data, num)

data, num = zip(*data_num)

Time Comparison

> for loop:
288 µs ± 12.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

> apply_along_axis Loop:
1.15 ms ± 43.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

> Compiled Loop:
5.45 µs ± 483 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

From the above comparison we see:

1. As @sai rightfully pointed out, NumPy's apply_along_axis function does not perform better than the regular loop, and my take on that is that due to the relatively small number of operations performed in each iteration, the speed gain achieved by using this function is dominated by the cost of its' header, so not only there is NO gain, but it actually slower than the regular loop.

2. The third way, i.e. the @jit performes the best, and it is because the for loop is compiled. So basically, now instead of python's interpreter, the just-in-time (or jit) compiler deals with this code now, which results in a c - like loop performance. More on Numba you could find here

Cheers.

This is a bit hacky but works:

>>> data = np.zeros((3, 4, 5))
>>> num = np.array([2, 4, 3])
>>> max_idx = min(data.shape[2], np.max(num))
>>> data[:,0,:max_idx] = np.arange(max_idx) < num[:,np.newaxis]
>>> data
array([[[1., 1., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.]],

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

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

It's based on the following idea:

>>> np.arange(5)
array([0, 1, 2, 3, 4])
>>> (np.arange(5) < 2).astype(np.float64)
array([1., 1., 0., 0., 0.])

except with 2 replaced by a vertically broadcasted num to generate the 'first N elements set to 1, rest 0'.

I'm not sure when using your exact data, but using an 3x4x5 example similar to yours, this code:

## make a 2d array of ones.
x = np.ones().reshape(3,1)

## mutiply those by an array of  [1,2,...,20]
## 20 because we want 4x5 max value
x = x*np.arange(1, 4*5+1)

## test where those are greater than the row number
## change the '3+1' expression to change the value of the comparison.
x = (x <= np.arange(1, 3+1).reshape(3, 1))

## booleans -> ints, reshape 2d to 3d (3x4x5) array
x = x.astype(int).reshape(-1, 4, 5)

produces this output:

array([[[1, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0]],

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

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

It's going to be a lot faster than iteration, so I hope it helps.

To make the code a little more generic,

## assign values for the axis lengths
ax_0_len = 3
ax_1_len = 4
ax_2_len = 5

## then replace the 3, 4, and 5 above with these variables.