Numpy高級索引用法

Question

案例1（解決）：數組A具有形狀（比如說）（300,50）。數組B是具有形狀（300,5）的索引數組，因此指示行i行i旁邊的「連續」行的索引。最終的結果是形狀爲（300,5,50）的數組C，例如C[i,j,:] = A[B[i,j],:]。這可以通過調用A[B,:]來完成。

這裏是殼體1小腳本實例：

import numpy as np 

## A is the data array 
A = np.arange(20).reshape((5,4)) 
## B indicate for each row which rows to pull together 
B = np.array([[0,2],[1,2],[2,0],[3,4],[4,1]]) 
A[B,:] #The desired result 

案例2（未解決）：同樣的問題，只是現在被成形（100,300,50）。如果B是形狀矩陣（100,300,5），則最終結果將是形狀爲（100,300,5,50）的陣列C，例如C[i,j,k,:] = A[i,B[i,j,k],:]。 A[B,:]不再工作，因爲它由於廣播而具有形狀（100,300,5,300,50）。

我該如何處理索引？

Answer 1

一種方法是重塑2D，保持列的數量不變，然後用扁平B索引將索引編入第一個軸，最後重新變形回到所需的值。

因此，實現起來 -

A.reshape(-1,A.shape[-1])[B.ravel()].reshape(100,300,5,50) 

那些重塑是僅僅觀看到陣列中，應該是相當有效的。

這解決了這兩種情況。這裏有一個樣品運行的情況下，1 -

1）輸入：

In [667]: A = np.random.rand(3,4) 
    ...: B = np.random.randint(0,3,(3,5)) 
    ...: 

2）原始方法：

In [668]: A[B,:] 
Out[668]: 
array([[[ 0.1 , 0.91, 0.1 , 0.98], 
     [ 0.1 , 0.91, 0.1 , 0.98], 
     [ 0.1 , 0.91, 0.1 , 0.98], 
     [ 0.45, 0.16, 0.02, 0.02], 
     [ 0.1 , 0.91, 0.1 , 0.98]], 

     [[ 0.45, 0.16, 0.02, 0.02], 
     [ 0.48, 0.6 , 0.96, 0.21], 
     [ 0.48, 0.6 , 0.96, 0.21], 
     [ 0.1 , 0.91, 0.1 , 0.98], 
     [ 0.45, 0.16, 0.02, 0.02]], 

     [[ 0.48, 0.6 , 0.96, 0.21], 
     [ 0.45, 0.16, 0.02, 0.02], 
     [ 0.48, 0.6 , 0.96, 0.21], 
     [ 0.45, 0.16, 0.02, 0.02], 
     [ 0.45, 0.16, 0.02, 0.02]]]) 

3）所提出的方法：

In [669]: A.reshape(-1,A.shape[-1])[B.ravel()].reshape(3,5,4) 
Out[669]: 
array([[[ 0.1 , 0.91, 0.1 , 0.98], 
     [ 0.1 , 0.91, 0.1 , 0.98], 
     [ 0.1 , 0.91, 0.1 , 0.98], 
     [ 0.45, 0.16, 0.02, 0.02], 
     [ 0.1 , 0.91, 0.1 , 0.98]], 

     [[ 0.45, 0.16, 0.02, 0.02], 
     [ 0.48, 0.6 , 0.96, 0.21], 
     [ 0.48, 0.6 , 0.96, 0.21], 
     [ 0.1 , 0.91, 0.1 , 0.98], 
     [ 0.45, 0.16, 0.02, 0.02]], 

     [[ 0.48, 0.6 , 0.96, 0.21], 
     [ 0.45, 0.16, 0.02, 0.02], 
     [ 0.48, 0.6 , 0.96, 0.21], 
     [ 0.45, 0.16, 0.02, 0.02], 
     [ 0.45, 0.16, 0.02, 0.02]]])