如何使用另一個 numpy 數組中的索引更新 numpy 數組

Question

我有

• 形狀為 (n1, n2, n3, n4) 的 numpy.array a
• 形狀為 (n1, n2, i1) 的索引數組idx

我想做的就像下面的代碼

for i in range(n1):
    for j in range(n2):
        for k in range(i1):
            b[i, j, k, :] = a[i, j, idx[i, j, k], :]

如果有一個 numpy function 來實現這個不用for循環？

Answer 1

使用作為起點：

import numpy as np

n1, n2, n3, n4, i1 = range(2, 7)

a = np.random.randint(10, size=(n1, n2, n3, n4))
idx = np.random.randint(n3, size=(n1, n2, i1))
b = np.zeros_like(a, shape=(n1, n2, i1, n4))

一般來說，您可以執行以下操作：

I, J, K = np.ogrid[:n1, :n2, :i1]
b[I, J, K] = a[I, J, idx]

這里的I J和K arrays 相當於循環變量i j和k 。 它們的形狀必須與idx的形狀一致。

如果b的形狀為 (n1, n2, i1, n4) 那么你不妨這樣做：

I, J, _ = np.ogrid[:n1, :n2, :1]
b = a[I, J, idx]

或者沒有ogrid ：

b = np.take_along_axis(a, idx[...,np.newaxis], axis=2)

這里newaxis用於插入一個長度為 1 的軸以允許廣播。 查看numpy 索引文檔以獲取更多信息。