If I have
x = np.arange(1, 10).reshape((3,3))
# array([[1, 2, 3],
# [4, 5, 6],
# [7, 8, 9]])
and
ind = np.array([[1,1], [1,2]])
# array([[1, 1],
# [1, 2]])
, how do I get use each row (axis 0) of ind to extract a cell of x? I hope to end up with the array [5, 6]. np.take(x, ind, axis=0) does not seem to work.
You could use "advanced integer indexing" by indexing x with two integer arrays, the first array for indexing the row, the second array for indexing the column:
In [58]: x[ind[:,0], ind[:,1]]
Out[58]: array([5, 6])
ind is actually huge. Is there any trick besides iteration?
ind[:,0] is an example of NumPy basic slicing, which returns a view. Views are very fast to create since the underlying data is not copied.x[ind.T.tolist()]
works, too, and can also be used for multidimensional NumPy arrays.
NumPy arrays are indexed by tuples. Usually, these tuples are created implicitly by python:
Note
In Python,
x[(exp1, exp2, ..., expN)]is equivalent tox[exp1, exp2, ..., expN];the latter is just syntactic sugar for the former.
Note that this syntactic sugar isn't NumPy-specific. You could use it on dictionaries when the key is a tuple:
In [1]: d = { 'I like the number': 1, ('pi', "isn't"): 2}
In [2]: d[('pi', "isn't")]
Out[2]: 2
In [3]: d['pi', "isn't"]
Out[3]: 2
Actually, it's not even related to indexing:
In [5]: 1, 2, 3
Out[5]: (1, 2, 3)
Thus, for your NumPy array, x = np.arange(1,10).reshape((3,3))
In [11]: x[1,2]
Out[11]: 6
because
In [12]: x[(1,2)]
Out[12]: 6
So, in unutbu's answer, actually a tuple containing the columns of ind is passed:
In [21]: x[(ind[:,0], ind[:,1])]
Out[21]: array([5, 6])
with x[ind[:,0], ind[:,1]] just being an equivalent (and recommended) short hand notation for the same.
Here's how that tuple looks like:
In [22]: (ind[:,0], ind[:,1])
Out[22]: (array([1, 1]), array([1, 2]))
We can construct the same tuple diffently from ind: tolist() returns a NumPy array's rows. Transposing switches rows and columns, so we can get a list of columns by first transposing and calling tolist on the result:
In [23]: ind.T.tolist()
Out[23]: [[1, 1], [1, 2]]
Because ind is symmetric in your example, it is it's own transpose. Thus, for illustration, let's use
In [24]: ind_2 = np.array([[1,1], [1,2], [0, 0]])
# array([[1, 1],
# [1, 2],
# [0, 0]])
In [25]: ind_2.T.tolist()
Out[25]: [[1, 1, 0], [1, 2, 0]]
This can easily be converted to the tuples we want:
In [27]: tuple(ind_2.T.tolist())
Out[27]: ([1, 1, 0], [1, 2, 0])
In [28]: tuple(ind.T.tolist())
Out[28]: ([1, 1], [1, 2])
Thus,
In [29]: x[tuple(ind.T.tolist())]
Out[29]: array([5, 6])
equivalently to unutbu's answer for x.ndim == 2 and ind_2.shape[1] == 2, but also working more generally when x.ndim == ind_2.shape[1], in case you have to work with multi-dimensional NumPy arrays.
Why you can drop the tuple(...) and directly use the list for indexing, I don't know. Must be a NumPy thing:
In [43]: x[ind_2.T.tolist()]
Out[43]: array([5, 6, 1])
FutureWarning: Using a non-tuple sequence for multidimensional indexing is deprecated; use arr[tuple(seq)]` instead of arr[seq]. In the future this will be interpreted as an array index, arr[np.array(seq)], which will result either in an error or a different result.` Do you know how to the same efficient indexing without the error?
tuple(...), so I'd try that.