NumPy數組使用argsort置換列3D矩陣

Question

我需要用從argsort獲得的2D置換矩陣pi置換矩陣A （ axis 0的3D矩陣）中列的元素，該矩陣包含所有列的新索引。

通過應用置換矩陣pi上的矩陣A （ A[pi]我將得到一個4D與新的形狀的矩陣。 例如，形狀A是（2,3,4）的形狀和A[pi]是（2,3,3,4）。

我能夠使用以下命令從A[pi]提取所需的排序矩陣：

swapaxes (diagonal(A[pi], axis1=2, axis2=1),1,2)

但這似乎太復雜和緩慢。

還有另一種優雅的解決方案嗎？

例：

print(A)
[[[   73   701  2411  2414]
  [ 5515  8292  8414 16135]
  [  100  1241  2146  2931]]

 [[ 1335  1747  3418  6312]
  [ 3788  5449  5753  9738]
  [  565  3038  3800  5430]]]

pi=argsort(Norm_order(A),0)

print(pi)
[[1, 0, 1],
 [0, 1, 0]]

print(swapaxes(diagonal(A[pi],axis1=2,axis2=1),1,2))
[[[ 1335  1747  3418  6312]
  [ 5515  8292  8414 16135]
  [  565  3038  3800  5430]]

 [[   73   701  2411  2414]
  [ 3788  5449  5753  9738]
  [  100  1241  2146  2931]]]

Answer 1

也許是一個口味問題，但是我發現以下內容更具可讀性：

i, j = np.ogrid[:3, :4]
A[pi[..., None], i, j]

輸出：

array([[[ 1335,  1747,  3418,  6312],
        [ 5515,  8292,  8414, 16135],
        [  565,  3038,  3800,  5430]],

       [[   73,   701,  2411,  2414],
        [ 3788,  5449,  5753,  9738],
        [  100,  1241,  2146,  2931]]])