numpy 高级索引技巧

numpy高级索引和索引技巧

NumPy提供了比常规Python序列更多的索引功能。如前所述，除了通过整数和切片建立索引外，还可以通过整数数组和布尔数组来建立数组索引。

用索引数组索引

>>> a = np.arange(12)**2                       # the first 12 square numbers
>>> 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])
>>>
>>> j = np.array([[3, 4], [9, 7]])      # a bidimensional array of indices
>>> a[j]                                       # the same shape as j
array([[ 9, 16],
       [81, 49]])

当索引数组a为多维时，单个索引数组引用的第一维a。下面的示例通过使用调色板将标签图像转换为彩色图像来显示此行为。

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

我们还可以为多个维度提供索引。每个维度的索引数组必须具有相同的形状。

>>> 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放在一个序列中（例如一个列表），然后对该列表进行索引。

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

但是，我们不能通过将i其j放入数组来完成此操作，因为该数组将被解释为索引a的第一个维度。

>>> s = np.array([i, j])

# not what we want
>>> a[s]
Traceback (most recent call last):
  File "<stdin>", line 1, in ?
IndexError: index (3) out of range (0<=index<=2) in dimension 0

# same as a[i, j]
>>> a[tuple(s)]
array([[ 2,  5]