I have a multidimensional numpy array and I would like to take some of its elements to construct a one dimensional array. The elements that I need to take are given by their indices, for example:
inds = [(0,0), (0,1), (1,1), (1,0), (0,2)]
I solve it in a straightforward way:
ls = []
for i, j in inds:
ls += [a[i,j]]
It gives the desired result. However, I have realized that this solution is too slow for my purposes. Are there a possibility to do the same in a more efficient way?
numpy arrays can be indexed with sequences (and, more generally, numpy arrays).
For example, here's my array a
In [19]: a
Out[19]:
array([[ 0, 1, 2, 3, 4],
[ 5, 6, 7, 8, 9],
[10, 11, 12, 13, 14]])
i and j hold the sequences of the first and second coordinates of your inds array:
In [20]: i
Out[20]: [0, 0, 1, 1, 0]
In [21]: j
Out[21]: [0, 1, 1, 0, 2]
You can use these to pull the corresponding values out of a:
In [22]: a[i, j]
Out[22]: array([0, 1, 6, 5, 2])
If you already have inds in your code, you can separate the list of tuples into i and j using zip:
In [23]: inds
Out[23]: [(0, 0), (0, 1), (1, 1), (1, 0), (0, 2)]
In [24]: i, j = zip(*inds)
In [25]: i
Out[25]: (0, 0, 1, 1, 0)
In [26]: j
Out[26]: (0, 1, 1, 0, 2)
Or, if inds is an array with shape (n, 2), like so:
In [27]: inds = np.array(inds)
In [28]: inds
Out[28]:
array([[0, 0],
[0, 1],
[1, 1],
[1, 0],
[0, 2]])
you can simply assign the transpose of inds to i, j:
In [33]: i, j = inds.T
In [34]: i
Out[34]: array([0, 0, 1, 1, 0])
In [35]: j
Out[35]: array([0, 1, 1, 0, 2])
In [36]: a[i, j]
Out[36]: array([0, 1, 6, 5, 2])
