I have two NxMx3 arrays and I want to reduce over the last dimension of the first array by selecting those elements corresponding to the index of the maximum value of each 3-vector of the second array to give an NxM result.

Hopefully that makes sense? If not hopefully the example below will shed some light.

Can anyone think of a more efficient way to do this than my 1st attempt? I thought that nditer might be the solution; I got it to work (mostly by trial & error) but found that it's ~50x slower for this task! Is this not a good usecase for nditer or am I doing something wrong?

In [42]: value_array = np.outer(ones(1000), arange(3)).reshape(20,50,3) ...: index_array = randn(*value_array.shape)

In [43]: indices = index_array.reshape(-1,3).argmax(axis=1) ...: result = value_array.reshape(-1,3)[np.arange(indices.size), indices] ...: result = result.reshape(value_array.shape[0:-1])

In [44]: it = np.nditer([value_array, index_array, None], ...: flags=['reduce_ok', 'external_loop','buffered', 'delay_bufalloc'], ...: op_flags=[['readonly'],['readonly'],['readwrite', 'allocate']], ...: op_axes=[[0,1,2], [0,1,2], [0,1,-1]]) ...: it.reset() ...: for values, index_values, out in it: ...: out[...] = values[index_values.argmax()] ...: #

In [45]: allclose(result, it.operands[2]) Out[45]: True

In [46]: %%timeit ...: indices = index_array.reshape(-1,3).argmax(axis=1) ...: result = value_array.reshape(-1,3)[np.arange(indices.size), indices] ...: result = result.reshape(value_array.shape[0:-1]) ...: 10000 loops, best of 3: 113 µs per loop

In [47]: %%timeit ...: it = np.nditer([value_array, index_array, None], ...: flags=['reduce_ok', 'external_loop','buffered', 'delay_bufalloc'], ...: op_flags=[['readonly'],['readonly'],['readwrite', 'allocate']], ...: op_axes=[[0,1,2], [0,1,2], [0,1,-1]]) ...: it.reset() ...: for values, index_values, out in it: ...: out[...] = values[index_values.argmax()] ...: # ...: 100 loops, best of 3: 5.26 ms per loop

In [48]:

Thanks, Dave