I have an array with shape (n, 2, 3) as:
array = np.array([[[-0.903, -3.47, -0.946], [-0.883, -3.48, -0.947]],
[[-1.02, -3.45, -0.992], [-1.01, -3.46, -1]],
[[-1.02, -3.45, -0.992], [-0.998, -3.45, -1]],
[[-0.638, -3.5, -0.897], [-0.604, -3.51, -0.896]],
[[-0.596, -3.52, -0.896], [-0.604, -3.51, -0.896]]])
and an index array for the second axis in which each value refer to each of two combinations e.g. for [-0.903, -3.47, -0.946], [-0.883, -3.48, -0.947] if the corresponding value in index array be 1, [-0.883, -3.48, -0.947] must be taken:
indices = np.array([0, 1, 0, 0, 1], dtype=np.int64)
the resulted array must be as below with shape (n, 3):
[-0.903, -3.47, -0.946] [-1.01, -3.46, -1] [-1.02, -3.45, -0.992] [-0.638, -3.5, -0.897] [-0.604, -3.51, -0.896]
How could I do so on a specified dimension just by NumPy.
>Solution :
In numpy you can combine slices along two dimensions. If you do arr[idx_x, idx_y] where idx_x and idx_y are 1d arrays of the same length you will get array of elements: [arr[idx_x[0], idx_y[0]], arr[idx_x[1], idx_y[1]], arr[idx_x[2], idx_y[2]], ...]
In your example if you do:
indices = np.array([0, 1, 0, 0, 1], dtype=np.int64)
x_idxs = np.arange(len(indices), dtype=int)
print(array[x_idxs, indices])
This will return result you want.