2.3.3 一个复杂的例子

下面让我们用所学的下标存取的知识，解决在NumPy邮件列表中提出的一个比较经典的问题。

我们对问题进行一些简化，提问者想要实现的下标运算是：有一个形状为(I, J, K)的三维数组v和一个形状为(I, J)的二维数组idx，idx的每个值都是0到K-L的整数。他想通过下标运算得到一个数组r，对于第0轴和第1轴的每个下标i和j都满足下面的条件：

    r[i,j,:] = v[i,j,idx[i,j]:idx[i,j]+L]

如图2-7所示，下左图中不透明的方块是我们希望获取的部分，通过下标运算之后将得到右侧所示的数组。

图2-7 三维数组下标运算问题的示意图

首先创建一个方便调试的数组v，它在第2轴上每一层的值就是该层的高度，即v[:,:,i]的所有的元素值都为i。然后随机产生数组idx，它的每个元素的取值都在0到K-L之间：

    I, J, K, L = 6, 7, 8, 3
    _, _, v = np.mgrid[:I, :J, :K]
    idx = np.random.randint(0, K - L, size=(I, J))

然后用数组idx创建第2轴的下标数组idx_k，它是一个形状为(I,J,L)的三维数组。它的第2轴上的每一层的值都等于idx数组加上层的高度，即idx_k[:,:,i]=idx[:,:] + i：

    idx_k = idx[:, :, None] + np.arange(3)
    idx_k.shape
    (6, 7, 3)

然后分别创建第0轴和第1轴的下标数组，它们的shape分别为(I,1,1)和(1,J,1)：

    idx_i, idx_j, _ = np.ogrid[:I, :J, :K]

使用idx_i, idx_j, idx_k对数组v进行下标运算即可得到结果：

    r = v[idx_i, idx_j, idx_k]    
    i, j = 2, 3  # 验证结果，读者可以将之修改为使用循环验证所有的元素
    r[i,j,:]  v[i,j,idx[i,j]:idx[i,j]+L]
    ---------  --------------------------
    [0, 1, 2]  [0, 1, 2]