Numpy：使用二維索引數組訪問二維數組

Question

我有兩個數組，一個是索引對矩陣，

a = array([[[0,0],[1,1]],[[2,0],[2,1]]], dtype=int)

另一個是在這些索引處訪問的數據矩陣

b = array([[1,2,3],[4,5,6],[7,8,9]])

我希望能夠使用 a 的索引來獲取b的條目。 只是做：

>>> b[a]

不工作，因為它給出了B的一行中的每個條目a ，即

array([[[[1,2,3],
         [1,2,3]],

        [[4,5,6],
         [4,5,6]]],


       [[[7,8,9],
         [1,2,3]],

        [[7,8,9],
         [4,5,6]]]])

當我想使用a的最后一個軸中的索引對來給出b的兩個索引時：

array([[1,5],[7,8]])

是否有這樣做的一個干凈的方式，或者我需要重塑b和組合的列a以相應的方式？

在我的實際問題中， a大約有 500 萬個條目， b是 100×100，我想避免 for 循環。

Answer 1

實際上，這是有效的：

b[a[:, :, 0],a[:, :, 1]]

給出array([[1, 5], [7, 8]]) 。

Answer 2

對於這種情況，這有效

tmp =  a.reshape(-1,2)
b[tmp[:,0], tmp[:,1]] 

Answer 3

一個更通用的解決方案，每當您想要使用具有任意大維度 m的形狀 (n,m) 索引的二維數組，命名為inds ，以便訪問另一個二維形狀數組 (n,k) 的元素，命名為B ：

# array of index offsets to be added to each row of inds
offset = np.arange(0, inds.size, inds.shape[1])

# numpy.take(B, C) "flattens" arrays B and C and selects elements from B based on indices in C
Result = np.take(B, offset[:,np.newaxis]+inds)

另一個不使用np.take並且我覺得更直觀的解決方案如下：

B[np.expand_dims(np.arange(B.shape[0]), -1), inds]

這種語法的優點是它既可以用於基於inds （如np.take ）從B讀取元素，也可以用於賦值。

您可以使用以下方法進行測試，例如：

B = 1/(np.arange(n*m).reshape(n,-1) + 1)
inds = np.random.randint(0,B.shape[1],(B.shape[0],B.shape[1]))