多维数组的下标存取

在前面的介绍中,我们通过一些实例介绍了如何对多维数组进行下标访问。实际上NumPy所提供的下标功能十分强大,在读者掌握了“广播”相关的知识之后,让我们再回过头来系统地学习数组的下标规则。

下标对象

首先多维数组的下标应该是一个长度和数组的维数相同的元组。如果下标元组的长度比数组的维数大则出错,如果小则在下标元组的后面补“:”,使得它的长度与数组维数相同。

如果下标对象不是元组,则NumPy会首先把它转换为元组。这种转换可能会和用户所希望的不一致,因此为了避免出现问题,请显式地使用元组作为下标。例如数组a是一个三维数组,下面使用一个二维列表lidx和二维数组aidx作为下标,所得到的结果不一样。

>>> a = np.arange(3*4*5).reshape(3,4,5)
>>> lidx = [[0],[1],[2]]
>>> aidx = np.array(lidx)
>>> a[lidx]
array([7])
>>> a[aidx]
array([[[[ 0,  1,  2,  3,  4],
         [ 5,  6,  7,  8,  9],
[[省略]]

这是因为NumPy将列表lidx转换成了([0],[1],[2]),而将数组aidx转换成了(aidx,:,:):

>>> a[tuple(lidx)]
array([7])
>>> a[aidx,:,:]
array([[[[ 0,  1,  2,  3,  4],
         [ 5,  6,  7,  8,  9],
[[省略]]

经过各种转换和添加“:”之后得到了一个标准的下标元组。它的各个元素有几种类型:切片、整数、整数数组和布尔数组。如果元素不是这些类型,如列表或元组,则将其转换成整数数组。

如果下标元组的所有元素都是切片和整数,那么用它作为下标得到的是原始数组的一个视图,即它和原始数组共享数据存储空间。

整数数组作为下标

下面让我们看看下标元组中的元素由切片和整数数组构成的情况。假设整数数组有N_t个,而切片有N_s个。N_t+N_s为数组的维数D

首先这N_t个整数数组必须满足广播条件,假设它们进行广播之后的维数为M,其形状为(d_0, d_1, ..., d_{M-1})

如果N_s为0 ,即没有切片元素时,则下标所得到的结果数组result的形状和整数数组广播之后的形状相同。它的每个元素值按照下面的公式获得:

result[i_0, i_1, ..., i_{M-1}] = X[ind_0[i_0, i_1, ..., i_{M-1}], ..., ind_{N_t-1}[i_0, i_1, ..., i_{M-1}]]

其中ind_0ind_{N_t-1}为进行广播之后的整数数组。让我们看一个例子加深对此公式的理解:

>>> i0 = np.array([[1,2,1],[0,1,0]])
>>> i1 = np.array([[[0]],[[1]]])
>>> i2 = np.array([[[2,3,2]]])
>>> b = a[i0, i1, i2]
>>> b
array([[[22, 43, 22],
        [ 2, 23,  2]],
       [[27, 48, 27],
        [ 7, 28,  7]]])

首先i0、i1、i2三个整数数组的shape属性分别为(2,3)、(2,1,1)、(1,1,3),根据广播规则,先在长度不足3的shape属性前面补1,使得它们的维数相同,广播之后的shape属性为各个轴的最大值:

(1, 2, 3)
(2, 1, 1)
(1, 1, 3)
---------
 2  2  3

即三个整数数组广播之后的shape属性为(2,2,3),这也就是下标运算所得到的结果数组的维数:

>>> b.shape
(2, 2, 3)

我们可以使用broadcast_arrays()查看广播之后的数组:

>>> ind0, ind1, ind2 = np.broadcast_arrays(i0, i1, i2)
>>> ind0
array([[[1, 2, 1],
        [0, 1, 0]],
       [[1, 2, 1],
        [0, 1, 0]]])
>>> ind1
array([[[0, 0, 0],
        [0, 0, 0]],
       [[1, 1, 1],
        [1, 1, 1]]])
>>> ind2
array([[[2, 3, 2],
        [2, 3, 2]],
       [[2, 3, 2],
        [2, 3, 2]]])

对于b中的任意一个元素b[i,j,k],它是数组a中经过ind0、ind1和ind2进行下标转换之后的值:

>>> i,j,k = 0,1,2
>>> b[i,j,k]
2
>>> a[ind0[i,j,k],ind1[i,j,k],ind2[i,j,k]]
2
>>> i,j,k = 1,1,1
>>> b[i,j,k]
28
>>> a[ind0[i,j,k],ind1[i,j,k],ind2[i,j,k]]
28

下面考虑N_s不为0的情况。当存在切片下标时,情况就变得更加复杂了。可以细分为两种情况:下标元组中的整数数组之间没有切片,即整数数组只有一个或者连续的。这时结果数组的shape属性为:将原始数组的shape属性中整数数组所占据的部分替换为它们广播之后的shape属性。例如假设原始数组a的shape属性为(3,4,5),i0和i1广播之后的形状为(2,2,3),则a[1:3,i0,i1]的形状为(2,2,2,3):

>>> c=a[1:3, i0, i1]
>>> c.shape
(2, 2, 2, 3)

其中c的shape属性中的第一个2是切片“1:3”的长度,后面的(2,2,3)则是i0和i1广播之后的数组的形状:

>>> ind0, ind1 = np.broadcast_arrays(i0, i1)
>>> ind0.shape
    (2, 2, 3)
>>> i,j,k = 1,1,2
>>> c[:,i,j,k]
    array([21, 41])
>>> a[1:3,ind0[i,j,k],ind1[i,j,k]] # 和c[:,i,j,k]的值相同
    array([21, 41])

当下标元组中的整数数组不是连续时,结果数组的shape属性为整数数组广播之后的形状后面添加上切片元素所对应的形状。例如a[i0,:,i1]的shape属性为(2,2,3,4)。其中(2,2,3)是i0和i1广播之后的形状,而4是数组a第1轴的长度:

>>> d = a[i0, :, i1]
>>> d.shape
    (2, 2, 3, 4)
>>> i,j,k = 1,1,2
>>> d[i,j,k,:]
    array([ 1,  6, 11, 16])
>>> a[ind0[i,j,k],:,ind1[i,j,k]]
    array([ 1,  6, 11, 16])

一个复杂的例子

下面让我们用所学的下标存取的知识,解决在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]

【图:三维数组下标运算问题的示意图】所示[1],左图中不透明的方块是我们希望获取的部分,通过下标运算之后将得到如右图所示的数组。

/tech/static/books/scipy/_images//numpy_array_index.png

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

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

02-numpy/numpy_array_index.py

三维数组下标运算问题

>>> 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.reshape(I,J,1) + np.arange(L)
>>> 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,:]
array([4, 5, 6])
>>> v[i,j,idx[i,j]:idx[i,j]+L]
array([4, 5, 6])

Footnotes

[1]绘制此图的源程序为“numpy_array_index_demo.py”。

布尔数组作下标

当使用布尔数组直接作为下标对象或者元组下标对象中有布尔数组时,都相当于用nonzero()将布尔数组转换成一组整数数组,然后使用整数数组进行下标运算。

nonzeros(a)返回数组a中值不为零的元素的下标,它的返回值是一个长度为a.ndim(数组a的轴数)的元组,元组的每个元素都是一个整数数组,其值为非零元素的下标在对应轴上的值。例如对于一维布尔数组b1,nonzero(b1)所得到的是一个长度为1的元组,它表示b1[0]和b1[2]的值不为0(False)。

>>> b1 = np.array([True, False, True, False])
>>> np.nonzero(b1)
    (array([0, 2]),)

对于二维数组b2,nonzero(b2)所得到的是一个长度为2的元组。它的第0个元素是数组a中值不为0的元素的第0轴的下标,第1个元素则是第1轴的下标,因此从下面的结果可知b2[0,0]、b[0,2]和b2[1,0]的值不为0:

>>> b2 = np.array([[True, False, True], [True, False, False]])
>>> np.nonzero(b2)
    (array([0, 0, 1]), array([0, 2, 0]))

当布尔数组直接做为下标时,相当于使用由nonzero()转换之后的元组作为下标对象:

>>> a = np.arange(3*4*5).reshape(3,4,5)
>>> a[b2]
array([[ 0,  1,  2,  3,  4],
       [10, 11, 12, 13, 14],
       [20, 21, 22, 23, 24]])
>>> a[np.nonzero(b2)]
array([[ 0,  1,  2,  3,  4],
       [10, 11, 12, 13, 14],
       [20, 21, 22, 23, 24]])

当下标对象是元组,并且其中有布尔数组时,相当于将布尔数组展开为由nonzeros()转换之后的各个整数数组:

>>> a[1:3, b2]
array([[20, 22, 25],
       [40, 42, 45]])
>>> a[1:3, np.nonzero(b2)[0], np.nonzero(b2)[1]]
array([[20, 22, 25],
       [40, 42, 45]])

內容目录

上一个主题

ufunc运算

下一个主题

庞大的函数库

本页

loading...