简体   繁体   中英

Indexing a multi-dimensional array partially

import numpy as np
arr = np.random.rand(50,3,3,3,16)
ids = (0,0,2,10)
b = arr[:, ids]  # don't work
b = arr[:, *ids]  # don't work
b = arr[:][ids]  # don't work
b = arr[:, tuple(ids)]  # don't work
b = arr[: + ids]  # don't work, obviously..
# b = arr[:,0,0,2,10].shape  # works (desired outcome)

I know there have been several questions about this, like Tuple as index of multidimensional array or Unpacking tuples/arrays/lists as indices for Numpy Arrays but none of them work for my case. Basically I want to index everything in the first axis, of specified "columns" in the rest of the axes (see the last line of my code). The desired output shape should be (50,) in this case.

But I want to index with a tuple/list of ids because I need to iterate through them, for example:

all_ids = ((0,0,0,2), (0,0,0,6), (1,1,0,2), (1,1,0,6),
           (2,2,0,2), (2,2,0,6), (2,2,2,2), (2,2,2,6))
c = 0
for id in all_ids:
    c += arr[:, id].sum() 

Add slice(None) to first dimension in ids and then subset:

arr[(slice(None),) + ids].shape
# (50,)

where:

(slice(None),) + ids
# (slice(None, None, None), 0, 0, 2, 10)

Notice slice(None, None, None) is equivalent to : , ie slice all. You can read docs on using slice object for indexing here .

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