loganecolss - 1 year ago 80
Python Question

# using an numpy array as indices of the 2nd dim of another array?

For example, I have two numpy arrays,

``````A = np.array(
[[0,1],
[2,3],
[4,5]])
B = np.array(
[[1],
[0],
[1]], dtype='int')
``````

and I want to extract one element from each row of
`A`
, and that element is indexed by
`B`
, so I want the following results:

``````C = np.array(
[[1],
[2],
[5]])
``````

I tried
`A[:, B.ravel()]`
`B`
, not what I want. Also looked into
`np.take`
, seems not the right solution to my problem.

However, I could use
`np.choose`
by transposing
`A`
,

``````np.choose(B.ravel(), A.T)
``````

but any other better solution?

You can use `NumPy's purely integer array indexing` -

``````A[np.arange(A.shape[0]),B.ravel()]
``````

Sample run -

``````In [57]: A
Out[57]:
array([[0, 1],
[2, 3],
[4, 5]])

In [58]: B
Out[58]:
array([[1],
[0],
[1]])

In [59]: A[np.arange(A.shape[0]),B.ravel()]
Out[59]: array([1, 2, 5])
``````

Please note that if `B` is a `1D` array or a list of such column indices, you could simply skip the flattening operation with `.ravel()`.

Sample run -

``````In [186]: A
Out[186]:
array([[0, 1],
[2, 3],
[4, 5]])

In [187]: B
Out[187]: [1, 0, 1]

In [188]: A[np.arange(A.shape[0]),B]
Out[188]: array([1, 2, 5])
``````
Recommended from our users: Dynamic Network Monitoring from WhatsUp Gold from IPSwitch. Free Download