python - Numpy 高级索引用法

Question

情况 1(已解决):数组 A 的形状(例如)(300,50)。数组 B 是形状为 (300,5) 的索引数组，使得 B[i,j]表示行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(未解决):同样的问题，只是现在A的形状是(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

一种方法是 reshape 为2D，保持列数不变，然后使用展平的B索引索引到第一个轴，最后 reshape 回所需的轴.

因此，实现将是 -

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

那些仅仅是数组 View 的 reshape 应该非常有效。

这解决了这两种情况。这是案例 #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]]])