I have one 2d array:
>>> input = np.array([[1,2],[3,4]])
>>> input
array([[1, 2],
[3, 4]])
and I also have this array which has the same number of columns as the input array and each lines contain the index to extract from the matching column of the input array.
>>> indices = np.array([[0],[1]])
>>> indices
array([[0],
[1]])
In this example I would like to get the following array as output:
array([[1],
[4]])
Any way I can achieve that?
One way via NumPy array indexing.
A = np.array([[1,2],[3,4]])
idx = np.array([[0],[1]])
res = A[np.arange(A.shape[0])[:, None], idx]
print(res)
[[1]
[4]]
Note we need to ensure both row and column indexers have the same shape, in this case (n, 1), where n is the number of rows.
A[idx, idx]?
A = np.array([[1,2],[3,4]]); idx = np.array([[1],[1]]); res = A[idx, idx] gives [4, 4]. Unless I've misread, OP expects [2, 4].You can just slice the index first:
in_array=np.arange(20).reshape(4,5)
ind=np.array([[0,2],[1,3]])
in_array[ind[:,0],ind[:,1]]
array([2, 8])
Also don't use the names such as input as they're built-in functions of python and you are overwriting them.
You can try this:
#!python
# -*- coding: utf-8 -*-#
#
# Imports
import numpy as np
arr = np.array([[1,2],[3,4]])
idx = np.array([[0],[1]])
ans = arr[idx,idx]
print(ans)
This answers the OP's specific question, however, if we want to make it more general answer, follow the idea of @jpp, and do like this:
# Imports
import numpy as np
arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
idx = np.array([0, 1, 1])
col_vec1 = np.arange(arr.shape[0])[:, None]
col_vec2 = idx[:, None]
ans = arr[col_vec1, col_vec2]
print(arr[0, 0]) # 1
print(arr[1, 1]) # 5
print(arr[2, 1]) # 8
print(ans)
