每列的numpy索引

Question

I want to get items from an array by indexing for each column. For example, an index like:

[[1, 3],
 [2, 4]]

Would mean:

1. get the 1st and 2nd element of the zeroth column of the array, and put them in the zeroth column of the result.

2. get the 3rd and 4th element of the 1st column of the array, and put them in the 1st column of the result.

Program:

import numpy as np
arr = np.arange(12).reshape(3,4)
print(arr)

ind = np.array(
    [[1,1,1,1],
     [2,2,2,2]]
)

result = arr[np.arange(2)[:, np.newaxis], ind[np.newaxis,:]]
print(result)

Result:

[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]

[[[1 1 1 1]
  [6 6 6 6]]]

Expected result:

[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]

[[[4 5 6 7]
  [8 9 10 11]]]

How to get the expected result?

Answer 1

This works:

import numpy as np
x = np.arange(12).reshape(3,4)

y = np.array(
    [[1,1,1,1],
     [2,2,2,2]]
)

z = x[y[:, np.newaxis], np.arange(4)[np.newaxis, :]]

print(z)

Result:

[[[ 4  5  6  7]]

 [[ 8  9 10 11]]]