简体   繁体   中英

Index a numpy array using a tuple (or list) when len < ndim?

I have a 3d numpy array, eg:

>>> A = np.arange(24).reshape(2,3,4)

I want to take a 1d slice along axis 0 based on a pair of coordinates for axes 1 and 2:

>>> h = 1
>>> l = 2
>>> A[:,h,l]
array([ 6, 18])

So far so good. But what if my coordinate pair is stored as a tuple or a list, rather than two integers? I've experimented with a few obvious options, to no avail:

>>> coords = (1,2)
>>> A[coords]
array([20, 21, 22, 23])
>>> A[:,coords]
array([[[ 4,  5,  6,  7],
        [ 8,  9, 10, 11]],

       [[16, 17, 18, 19],
        [20, 21, 22, 23]]])
>>> A[...,coords]
array([[[ 1,  2],
        [ 5,  6],
        [ 9, 10]],

       [[13, 14],
        [17, 18],
        [21, 22]]])

I've googled around on this and not found anything, but it's entirely possible that I'm not searching with the appropriate jargon. So, apologies if this is an overly simplistic question!

You can construct the slice tuple directly, with something like:

In [11]: A[(slice(None),) + coords]
Out[11]: array([ 6, 18])

This is because calling A[:, 1, 2] is equivalent / calls:

In [12]: A.__getitem__((slice(None, None, None), 1, 2))
Out[12]: array([ 6, 18])

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM