Index n dimensional array with (n-1) d array

Question

What is the most elegant way to access an n-dimensional array with an (n-1) -dimensional array at a given size, as in the dummy example

a = np.random.random_sample((3,4,4)) b = np.random.random_sample((3,4,4)) idx = np.argmax(a, axis=0) 

How can I access now using idx a to get the maximum values ​​in a , as if I were using a.max(axis=0) ? or how to get the values ​​specified by idx in b ?

I was thinking about using np.meshgrid , but I think this is overkill. Note that the axis dimension can be any usable axis (0,1,2) and is not known in advance. Is there an elegant way to do this?

Answer 1

Use advanced-indexing -

 m,n = a.shape[1:] I,J = np.ogrid[:m,:n] a_max_values = a[idx, I, J] b_max_values = b[idx, I, J] 

---

In general:

 def argmax_to_max(arr, argmax, axis): """argmax_to_max(arr, arr.argmax(axis), axis) == arr.max(axis)""" new_shape = list(arr.shape) del new_shape[axis] grid = np.ogrid[tuple(map(slice, new_shape))] grid.insert(axis, argmax) return arr[tuple(grid)] 

Pretty little uncomfortable than such a natural operation should be, unfortunately.

To index an n dim array with an (n-1) dim array, we could simplify it a bit to give us a grid of indices for all axes, for example:

 def all_idx(idx, axis): grid = np.ogrid[tuple(map(slice, idx.shape))] grid.insert(axis, idx) return tuple(grid) 

Therefore, use it to index into input arrays -

 axis = 0 a_max_values = a[all_idx(idx, axis=axis)] b_max_values = b[all_idx(idx, axis=axis)]