Is there an efficient way to index over multiple arrays?
For example, I have an array I want to index from
a = [[1,2,3],[4,5,6]]
And another array contains the indices.
b = [[0, 1], [1,2]]
And I expect [[1, 2], [5, 6]], which indexes the first row of a by [0,1], and indexes the second row of a by [1,2].
Thanks.
If a and b are of same length, may be you can try using np.take as following:
import numpy as np
a = [[1,2,3],[4,5,6]]
b = [[0, 1], [1,2]]
result = [np.take(a[i],b[i]).tolist() for i in range(len(a))]
print(result)
# result: [[1, 2], [5, 6]]
In [107]: a = [[1,2,3],[4,5,6]]
In [108]: b = [[0, 1], [1,2]]
a and b are lists. The appropriate solution is a nested list comprehension
In [111]: [[a[i][j] for j in x] for i,x in enumerate(b)]
Out[111]: [[1, 2], [5, 6]]
Now if a is made into a numpy array:
In [112]: np.array(a)[np.arange(2)[:,None], b]
Out[112]:
array([[1, 2],
[5, 6]])
For this the 1st dimension of the array is indexed with a (2,1) array, and the 2nd with a (2,2). They broadcast together to produce a (2,2) result.
is working in the same direction, but the accepted answer use ix_
Y[np.ix_([0,3],[0,3])]
which won't work in the case of a (2,2) b.
In [113]: np.array(a)[np.ix_(np.arange(2), b)]
ValueError: Cross index must be 1 dimensional
ix_ will turn the 1st dimension np.arange(2) in to the right (2,1).
This might make the broadcasting more explicit:
In [114]: np.array(a)[[[0,0],[1,1]], [[0,1],[1,2]]]
Out[114]:
array([[1, 2],
[5, 6]])
It selects elements (0,0), (0,1), (1,1) and (1,2)
To further test this, make b non symmetic:
In [138]: b = [[0, 1,1], [1,2,0]] # (2,3)
In [139]: np.array(a)[np.arange(2)[:,None], b]
Out[139]:
array([[1, 2, 2],
[5, 6, 4]])
