部分索引多維數組

Question

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)

我知道對此有幾個問題，例如元組作為多維數組的索引，或者拆包元組/數組/列表作為Numpy數組的索引，但是它們都不適合我的情況。 基本上，我想在第一個軸上索引所有其余軸中指定的“列”中的所有內容（請參閱代碼的最后一行）。 在這種情況下(50,)所需的輸出形狀應為(50,) 。

但是我想用元組/ ID列表建立索引，因為我需要遍歷它們，例如：

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() 

Answer 1

將slice(None)添加到ids第一維，然后添加子集：

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

哪里：

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

注意slice(None, None, None)等效於: ，即全部切片。 您可以在此處閱讀有關使用slice對象進行索引的文檔 。