Python 使用索引元組切片 numpy 數組

Question

我想使用 N-1 維索引元組b對 N 維 numpy 數組a進行切片。 下面是 N = 4 的情況：

import numpy as np

a = np.random.rand(20,1,5,4)
b1 = (0,0,0)
b2 = (0,3,2)

# want to get the slice c = a[:,b2[0],b2[1],b2[2]] without explicitly writing all N-1 elements of b2

c = a[:,b1] # no error message, but gives a[:,[b1[0],b1[1],b1[2]],:,:]
c = a[:,b2] # error message (because b2[1] = 3 larger than a.shape[1])

將元組索引元素“解包”到不同維度的正確方法是什么？ 我嘗試使用星號運算符 ( *b )，但這似乎不起作用。

Answer 1

a[(..., *b1)]

這歸結為使用索引元組中的變量參數手動調用__getitem__方法：

In [23]: a.__getitem__((..., *b1))
Out[23]: 
array([0.93049554, 0.63858628, 0.31665933, 0.15048072, 0.17083849,
       0.34621308, 0.66980288, 0.75191251, 0.05844084, 0.16715787,
       0.80409005, 0.64181912, 0.76463963, 0.75616746, 0.44955961,
       0.19034925, 0.8841084 , 0.54450091, 0.02033021, 0.21275405])

In [21]: a[:, 0, 0, 0]
Out[21]: 
array([0.93049554, 0.63858628, 0.31665933, 0.15048072, 0.17083849,
       0.34621308, 0.66980288, 0.75191251, 0.05844084, 0.16715787,
       0.80409005, 0.64181912, 0.76463963, 0.75616746, 0.44955961,
       0.19034925, 0.8841084 , 0.54450091, 0.02033021, 0.21275405])

請注意，省略號 ( ... ) 的使用具有更一般的含義。 否則，您必須手動構造:切片：

In [20]: a.__getitem__((slice(None), *b1))
Out[20]: 
array([0.93049554, 0.63858628, 0.31665933, 0.15048072, 0.17083849,
       0.34621308, 0.66980288, 0.75191251, 0.05844084, 0.16715787,
       0.80409005, 0.64181912, 0.76463963, 0.75616746, 0.44955961,
       0.19034925, 0.8841084 , 0.54450091, 0.02033021, 0.21275405])