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]]])