How to pick one of two arrays in an axis of multidimensional NumPy array with an 1D index array for-CodePudding

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.

CodePudding user response：

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.

CodePudding user response：

Try with a for loop:

out = []
for i in range(len(indices)):
    out.append(list(arr[i,indices[i]]))

print(out)

Output:
[[-0.903, -3.47, -0.946],
 [-1.01, -3.46, -1.0],
 [-1.02, -3.45, -0.992],
 [-0.638, -3.5, -0.897],
 [-0.604, -3.51, -0.896]]