NumPy之 索引技巧

系列文章

一次性搞定NumPy入门基础知识
NumPy之操控ndarray的形状
NumPy之浅拷贝和深拷贝
NumPy之索引技巧

利用array作为索引

索引一维ndarray

当被索引的ndarray是一维时，利用array做索引，相当于一次性从被索引对象中挑选出索引指定的所有元素，索引出的对象仍然是一个ndarray对象。

>>> a = np.arange(12)**2     
>>> a
array([  0,   1,   4,   9,  16,  25,  36,  49,  64,  81, 100, 121])
>>> i = np.array( [ 1,1,3,8,5 ] )              # an array of indices
>>> a[i]                                       # the elements of a at the positions i
array([ 1,  1,  9, 64, 25])

如果索引是多维的，那么检出的ndarray的shape和索引的形状一致，索引中的各个数字则分别代表了元素在被检索ndarray中的位置

>>> a = np.arange(12)**2 
>>> a
array([  0,   1,   4,   9,  16,  25,  36,  49,  64,  81, 100, 121])
>>> j = np.array( [ [ 3, 4], [ 9, 7 ] ] )      # a bidimensional array of indices
>>> a[j]                                       # the same shape as j
array([[ 9, 16],
       [81, 49]])

索引多维ndarray

当被索引对象是多维的，且只提供一个索引，那么索引中的每个数字所对应的元素是被索引对象的第一个维度。

>>> palette = np.array( [ [0,0,0],                # black
...                       [255,0,0],              # red
...                       [0,255,0],              # green
...                       [0,0,255],              # blue
...                       [255,255,255] ] )       # white
>>> image = np.array( [ [ 0, 1, 2, 0 ],           # each value corresponds to a color in the palette
...                     [ 0, 3, 4, 0 ]  ] )
>>> palette[image]                            # the (2,4,3) color image
array([[[  0,   0,   0],
        [255,   0,   0],
        [  0, 255,   0],
        [  0,   0,   0]],
       [[  0,   0,   0],
        [  0,   0, 255],
        [255, 255, 255],
        [  0,   0,   0]]])

如果被索引对象是多维的，并且想从被索引对象中进行多维检索，那么需要给出多组索引array，这个时候每个索引array对应被索引对象中的一个维度。需要注意，索引array的形状必须是一致的，因为必须使用这几个索引array的相同位置（构成了被索引对象的不同维度）来共同指定被索引对象中的某一个元素。

>>> a = np.arange(12).reshape(3,4)
>>> a
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
>>> i = np.array( [ [0,1],                        # indices for the first dim of a
...                 [1,2] ] )
>>> j = np.array( [ [2,1],                        # indices for the second dim
...                 [3,3] ] )
>>>
>>> a[i,j]                                     # i and j must have equal shape
array([[ 2,  5],
       [ 7, 11]])
>>>
>>> a[i,2]
array([[ 2,  6],
       [ 6, 10]])
>>>
>>> a[:,j]                                     # i.e., a[ : , j]
array([[[ 2,  1],
        [ 3,  3]],
       [[ 6,  5],
        [ 7,  7]],
       [[10,  9],
        [11, 11]]])

通常可以把上述i和j写成list的形式，来作为索引使用

>>> l = [i,j]
>>> a[l]                                       # equivalent to a[i,j]
array([[ 2,  5],
       [ 7, 11]])

注意不能将i和j组合成array进行使用，这样会被解释器认为是索引被索引对象的第一个维度：

>>> s = np.array( [i,j] )
>>> a[s]                                       # not what we want
Traceback (most recent call last):
  File "<stdin>", line 1, in ?
IndexError: index (3) out of range (0<=index<=2) in dimension 0
>>>
>>> a[tuple(s)]                                # same as a[i,j]
array([[ 2,  5],
       [ 7, 11]])

下面的例子展示了多维数组索引的一个应用：

>>> time = np.linspace(20, 145, 5)                 # 设置一个时间序列
>>> data = np.sin(np.arange(20)).reshape(5,4)      # 对应于上述时间序列的4组数据
>>> time
array([  20.  ,   51.25,   82.5 ,  113.75,  145.  ])
>>> data
array([[ 0.        ,  0.84147098,  0.90929743,  0.14112001],
       [-0.7568025 , -0.95892427, -0.2794155 ,  0.6569866 ],
       [ 0.98935825,  0.41211849, -0.54402111, -0.99999021],
       [-0.53657292,  0.42016704,  0.99060736,  0.65028784],
       [-0.28790332, -0.96139749, -0.75098725,  0.14987721]])
>>>
>>> ind = data.argmax(axis=0)                  # 求出每组数据的最大值索引
>>> ind
array([2, 0, 3, 1])
>>>
>>> time_max = time[ind]                       # 检索出数据取最大值时的时间
>>>
>>> data_max = data[ind, range(data.shape[1])] # => # 检索出每组的最大数据，这里使用了二维索引，第一维是ind:[2,0,3,1]，第二维是数组的列号，分别是[0,1,2,3]。.
>>>
>>> time_max
array([  82.5 ,   20.  ,  113.75,   51.25])
>>> data_max
array([ 0.98935825,  0.84147098,  0.99060736,  0.6569866 ])
>>>
>>> np.all(data_max == data.max(axis=0)) #和系统函数data.max进行对比，检查检索是否正确
True

可以使用索引来进行赋值操作，但注意如果索引有重复，将以最后一次赋值为准：

>>> a = np.arange(5)
>>> a
array([0, 1, 2, 3, 4])
>>> a[[1,3,4]] = 0
>>> a
array([0, 0, 2, 0, 0])

>>> a = np.arange(5)
>>> a[[0,0,2]]=[1,2,3]
>>> a
array([2, 1, 3, 3, 4])

使用布尔值作为索引

上述内容使用整数作为索引，而当使用布尔值作为索引时，则是通过指定需要哪些内容，不需要哪些内容来实现的（True，False）。
下面的例子通过给出一个原始ndarray相同形状的索引来进行操作，非常直观：

>>> a = np.arange(12).reshape(3,4)
>>> b = a > 4                               # 通过逻辑运算得出一组布尔值
>>> b                                          # b的形状和一致
array([[False, False, False, False],
       [False,  True,  True,  True],
       [ True,  True,  True,  True]])
>>> a[b]                                       # 得出一个一维数组
array([ 5,  6,  7,  8,  9, 10, 11])

这个性质经常应用在赋值操作里：

>>> a[b] = 0             # 将所有大于4的元素设置为0                 
>>> a
array([[0, 1, 2, 3],
       [4, 0, 0, 0],
       [0, 0, 0, 0]])

另外一种用法类似于上述的多维度索引：对于每个维度，提供一个一维的array，说明需要这个维度上的哪些元素。

>>> a = np.arange(12).reshape(3,4)
>>> b1 = np.array([False,True,True])             # first dim selection
>>> b2 = np.array([True,False,True,False])       # second dim selection
>>>
>>> a[b1,:]                                   # selecting rows
array([[ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
>>>
>>> a[b1]                                     # same thing
array([[ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
>>>
>>> a[:,b2]                                   # selecting columns
array([[ 0,  2],
       [ 4,  6],
       [ 8, 10]])
>>>
>>> a[b1,b2]                                  # 相当于a[[1,2],[0,2]]
array([ 4, 10])

需要注意，这里b1和b2的长度要分别与a的两个维度的长度一致。

ix_()函数

ix()函数的参数是N个一维数组，它可以生产一个组合索引出来，这个索引有N维，每个维度的长度为每个参数的长度，取值由每个一维数组具体的值来决定。

>>> a = np.array([2,3,4,5])
>>> b = np.array([8,5,4])
>>> c = np.array([5,4,6,8,3])
>>> ax,bx,cx = np.ix_(a,b,c)
>>> ax
array([[[2]],
       [[3]],
       [[4]],
       [[5]]])
>>> bx
array([[[8],
        [5],
        [4]]])
>>> cx
array([[[5, 4, 6, 8, 3]]])
>>> ax.shape, bx.shape, cx.shape
((4, 1, 1), (1, 3, 1), (1, 1, 5))
>>> result = ax+bx*cx  # 这里用到broadcasting原则，可以暂不考虑
>>> result
array([[[42, 34, 50, 66, 26],
        [27, 22, 32, 42, 17],
        [22, 18, 26, 34, 14]],
       [[43, 35, 51, 67, 27],
        [28, 23, 33, 43, 18],
        [23, 19, 27, 35, 15]],
       [[44, 36, 52, 68, 28],
        [29, 24, 34, 44, 19],
        [24, 20, 28, 36, 16]],
       [[45, 37, 53, 69, 29],
        [30, 25, 35, 45, 20],
        [25, 21, 29, 37, 17]]])
>>> result[3,2,4]
17
>>> a[3]+b[2]*c[4]
17
>>> ix = np.ix_([0,1],[1,2],[1,2,3]) # 生成一个ix索引
>>> result[ix]    # 利用ix索引来生成新的ndarray
array([[[22, 32, 42],
        [18, 26, 34]],

       [[23, 33, 43],
        [19, 27, 35]]])