I have 2 numpy arrays: X with shape (..., 12, 3) and Y with shape (..., 12), where the ... is the same in both cases. I get an array of indices using:
idx = np.argmin(Y, axis=-1)
which afterwards I want to use to get the corresponding elements from array X and Y. However, I run into issues due to the extra dimensions. For Y, I can do e.g.
Y[np.arange(Y.shape[0]), idx]
however this requires me to know the shape of ... which will not be the case at all times. How do I do this in a better way?
Okay, since I figured it out, here is the way to do it:
idx = np.argmin(Y, axis=-1)
indices = np.indices(idx.shape)
min_from_X = X[indices[0], ..., idx, :]
min_from_Y = Y[indices[0], ..., idx]