給定索引列表，將任意維度的 numpy 切片到一維數組

Question

我有一個 numpy ndarray arr和indices ，一個指定特定條目的索引列表。 具體來說，讓我們采取：

arr = np.arange(2*3*4).reshape((2,3,4))
indices= [1,0,3]

我有代碼通過arr觀察除一個索引n之外的所有 1d 切片：

arr[:, indices[1], indices[2]]  # n = 0
arr[indices[0], :, indices[2]]  # n = 1
arr[indices[0], indices[1], :]  # n = 2

我想更改我的代碼以循環n並支持任意維度的arr 。

我查看了文檔中的 索引例程條目，並找到了有關slice()和np.s_()的信息。 我能夠拼湊出我想要的東西：

def make_custom_slice(n, indices):
    s = list()
    for i, idx in enumerate(indices):
        if i == n:
            s.append(slice(None))
        else:
            s.append(slice(idx, idx+1))
    return tuple(s)


for n in range(arr.ndim):
    np.squeeze(arr[make_custom_slice(n, indices)])

其中np.squeeze用於刪除長度為 1 的軸。沒有這個，這個產生的數組具有形狀(arr.shape[n],1,1,...)而不是(arr.shape[n],) .

有沒有更慣用的方法來完成這項任務？

Answer 1

對上述解決方案的一些改進（可能仍然存在單行或更高性能的解決方案）：

def make_custom_slice(n, indices):
    s = indices.copy()
    s[n] = slice(None)
    return tuple(s)


for n in range(arr.ndim):
    print(arr[make_custom_slice(n, indices)])

一個 integer 值idx可用於替換切片 object slice(idx, idx+1) 。 因為大多數索引都是直接復制的，所以從索引的副本開始，而不是從頭開始構建列表。

當以這種方式構建時， arr[make_custom_slice(n, indices)的結果具有預期的維度，並且np.squeeze是不必要的。