NumPy 高级索引

更新时间：2022-06-09 23:43:18 标签：numpy python 索引

好消息！

本教程作者所著新书《深入浅出Pandas：利用Python进行数据处理与分析》（ISBN：9787111685456）已由机械工业出版社出版上市，各大电商平台有售，欢迎：查看详情并关注购买。

说明

NumPy 教程正在计划编写中，欢迎大家加微信 gairuo123 提供意见、建议、纠错、催更。

当 NumPy 数组的选择对象 obj 是非元组序列对象、ndarray（数据类型为 integer 或 bool）或至少有一个序列对象或 ndarray（数据类型为integer或bool）的元组时，会触发高级索引。NumPy 有两种类型的高级索引：整型和布尔型。高级索引总是返回数据的副本（与基本索引切片返回视图的不同）。

与基本索引区别

高级索引的定义意味着 x[(1, 2, 3),] 与 x[(1, 2, 3)] 有根本不同。后者相当于 x[1, 2, 3]，它将触发基本选择，而前者将触发高级索引，一定要理解为什么会发生这种情况。

还要认识到，x[[1, 2, 3]] 将触发高级索引，而由于上述不推荐的数字兼容性，x[[1, 2, slice(None)]] 将触发基本切片。

整数数组索引

整数数组索引允许根据数组中的任意项的 N 维索引来选择它们。每个整数数组表示该维度中的多个索引。

索引数组中允许出现负值，并且与处理单个索引或切片时一样：

x = np.arange(10, 1, -1)
x
# array([10,  9,  8,  7,  6,  5,  4,  3,  2])
x[np.array([3, 3, 1, 8])]
# array([7, 7, 9, 2])
x[np.array([3, 3, -3, 8])]
# array([7, 7, 4, 2])

如果索引值超出界限，则引发 IndexError：

x = np.array([[1, 2], [3, 4], [5, 6]])
x[np.array([1, -1])]
'''
array([[3, 4],
      [5, 6]])
'''

x[np.array([3, 4])]
'''
Traceback (most recent call last):
  ...
IndexError: index 3 is out of bounds for axis 0 with size 3
'''

当索引由与被索引数组的维数一样多的整数数组组成时，索引很简单，但与切片不同。

高级索引始终作为一个索引进行广播和迭代：

result[i_1, ..., i_M] == x[ind_1[i_1, ..., i_M], ind_2[i_1, ..., i_M],
                           ..., ind_N[i_1, ..., i_M]]

请注意，生成的形状与（广播）索引数组形状 ind_1, ..., ind_N 相同。如果索引无法广播到同一形状，则异常索引器错误：IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes...。

使用多维索引数组进行索引往往不常用，但它们是允许的，并且对于某些问题很有用。我们将从最简单的多维情况开始：

y = np.arange(35).reshape(5, 7)
y
'''
array([[ 0,  1,  2,  3,  4,  5,  6],
       [ 7,  8,  9, 10, 11, 12, 13],
       [14, 15, 16, 17, 18, 19, 20],
       [21, 22, 23, 24, 25, 26, 27],
       [28, 29, 30, 31, 32, 33, 34]])
'''

y[np.array([0, 2, 4]), np.array([0, 1, 2])]
# array([ 0, 15, 30])

在这种情况下，如果索引数组具有匹配的形状，并且对于要索引的数组的每个维度都有一个索引数组，则生成的数组具有与索引数组相同的形状，并且值对应于索引数组中每个位置的索引集。在本例中，两个索引数组的第一个索引值都是0，因此结果数组的第一个值是 y[0，0]。下一个值是 y[2，1]，最后一个值是 y[4，2]。

如果索引数组不具有相同的形状，则会尝试将它们广播到相同的形状。如果无法将它们广播到同一形状，则会引发异常：

y[np.array([0, 2, 4]), np.array([0, 1])]
'''
Traceback (most recent call last):
  ...
IndexError: shape mismatch: indexing arrays could not be broadcast
together with shapes (3,) (2,)
'''

广播机制允许索引数组与其他索引的标量组合。其效果是标量值用于索引数组的所有对应值：

y[np.array([0, 2, 4]), 1]
# array([ 1, 15, 29])

跳转到下一个复杂级别，只可能使用索引数组对数组进行部分索引。要理解这种情况下会发生什么，需要一点思考。例如，如果我们只对y使用一个索引数组：

y[np.array([0, 2, 4])]
'''
array([[ 0,  1,  2,  3,  4,  5,  6],
      [14, 15, 16, 17, 18, 19, 20],
      [28, 29, 30, 31, 32, 33, 34]])
'''

它导致构建一个新数组，其中索引数组的每个值从被索引的数组中选择一行，结果数组具有结果形状（索引元素的数量、行的大小）。

通常，结果数组的形状将是索引数组的形状（或所有索引数组广播到的形状）与被索引数组中任何未使用维度（未索引维度）的形状的串联。

案例1

从每行中选择一个特定元素。行索引只是 [0，1，2]，列索引指定要为相应行选择的元素，这里是 [0，1，0]。将两者结合使用，可以使用高级索引解决该任务：

x = np.array([[1, 2], [3, 4], [5, 6]])
x
'''
array([[1, 2],
       [3, 4],
       [5, 6]])
'''

x[[0, 1, 2], [0, 1, 0]]
# array([1, 4, 5])

为了实现类似于上述基本切片的行为，可以使用广播。功能 ix_ 有助于此广播。最好通过一个例子来理解这一点。

案例2

从4x3阵列中，应使用高级索引选择角元素。因此，需要选择列为[0，2]之一、行为[0，3]之一的所有元素。要使用高级索引，需要显式选择所有元素。使用前面介绍的方法，可以写：

x = np.array([[ 0,  1,  2],
              [ 3,  4,  5],
              [ 6,  7,  8],
              [ 9, 10, 11]])

rows = np.array([[0, 0],
                 [3, 3]], dtype=np.intp)

columns = np.array([[0, 2],
                    [0, 2]], dtype=np.intp)
x[rows, columns]
'''
array([[ 0,  2],
       [ 9, 11]])
'''

然而，由于上面的索引数组只是重复它们自己，所以可以使用广播（比较操作，例如rows[:, np.newaxis] + columns)）来简化这一过程：

rows = np.array([0, 3], dtype=np.intp)
columns = np.array([0, 2], dtype=np.intp)
rows[:, np.newaxis]
'''
array([[0],
       [3]])
'''

x[rows[:, np.newaxis], columns]
'''
array([[ 0,  2],
       [ 9, 11]])
'''

此广播也可以使用功能 ix_ 来实现：

x[np.ix_(rows, columns)]
'''
array([[ 0,  2],
       [ 9, 11]])
'''

请注意，没有主调用 np.ix_ 仅选择对角线元素：

x[rows, columns]
# array([ 0, 11])

对于使用多个高级索引进行索引，这种差异是需要记住的最重要的一点。

案例3

高级索引可能有用的一个实际示例是颜色查找表，我们希望将图像的值映射到RGB三元组中进行显示。查找表可以有 shape (nlookup, 3)。使用 shape (ny, nx)且 dtype=np.uint8（或任何整数类型，只要值与查找表的边界一致）将生成一个形状数组shape (ny, nx, 3)，其中RGB值的三倍与每个像素位置相关联。

布尔数组索引

当obj是布尔类型的数组对象（可能从比较运算返回）时，会发生这种高级索引。单个布尔索引数组实际上与 x[obj.nonzero()] 相同。

其中，如上所述，obj.nonzero() 返回整数索引数组的元组（长度为 obj.ndim），显示 obj 中为 True 的元素。然而，当 obj.shape == x.shape 时速度更快。

当 obj.ndim == x.ndim 时，x[obj] 返回一个一维数组，其中填充与 obj 的真值相对应的 x 元素。搜索顺序将是 row-major, C-style。如果obj在x界限之外的条目上有真值，则会引发索引错误。如果obj 小于 x，则等于用 False 填充。

这方面的一个常见用例是过滤所需的元素值。例如，可能希望从非NaN的数组中选择所有条目：

x = np.array([[1., 2.], [np.nan, 3.], [np.nan, np.nan]])
x[~np.isnan(x)]
# array([1., 2., 3.])

或者希望为所有负元素添加一个常量：

x = np.array([1., -1., -2., 3])
x[x < 0] += 20
x
# array([ 1., 19., 18., 3.])

通常，如果索引包含布尔数组，则结果将与插入 obj.nonzero() 到相同的位置，并使用上面描述的整数数组索引机制。

x[ind_1, boolean_array, ind_2]
# 相当于
x[(ind_1,) + boolean_array.nonzero() + (ind_2,)]

如果只有一个布尔数组而没有整数索引数组，那么这很简单。只需注意确保布尔索引的维度与它应该使用的维度完全相同。

通常，当布尔数组的维数小于被索引的数组时，这相当于 x[b, ...]，这意味着x由b索引，然后是填写x的秩所需的 :。因此，结果的形状是一个包含布尔数组的真元素数的维度，然后是被索引数组的其余维度：

x = np.arange(35).reshape(5, 7)
b = x > 20
b[:, 5]
# array([False, False, False,  True,  True])
x[b[:, 5]]
'''
array([[21, 22, 23, 24, 25, 26, 27],
      [28, 29, 30, 31, 32, 33, 34]])
'''

这里，从索引数组中选择第4行和第5行，并将其组合成一个二维数组。

案例1

从数组中，选择总和小于或等于2的所有行：

x = np.array([[0, 1], [1, 1], [2, 2]])
rowsum = x.sum(-1)
x[rowsum <= 2, :]
'''
array([[0, 1],
       [1, 1]])
'''

组合多个布尔索引数组或布尔与整数索引数组，可以用 obj.nonzero() 类比理解。

函数 ix_ 还支持布尔数组，它可以正常工作。

案例2

使用布尔索引选择所有相加为偶数的行。同时，应使用高级整数索引选择列0和列2。使用 ix_ 功能，可以通过以下方式完成此操作：

x = np.array([[ 0,  1,  2],
              [ 3,  4,  5],
              [ 6,  7,  8],
              [ 9, 10, 11]])
rows = (x.sum(-1) % 2) == 0
rows
# array([False,  True, False,  True])
columns = [0, 2]
x[np.ix_(rows, columns)]
'''
array([[ 3,  5],
       [ 9, 11]])
'''

没有 np.ix_ 调用，仅选择对角线元素。或没有用 np.ix_（比较整数数组示例）：

rows = rows.nonzero()[0]
x[rows[:, np.newaxis], columns]
'''
array([[ 3,  5],
       [ 9, 11]])
'''

案例3

使用具有四个真元素的二维布尔形状数组（2，3），从三维形状数组（2，3，5）中选择行，结果为二维形状数组（4，5）：

x = np.arange(30).reshape(2, 3, 5)
x
'''
array([[[ 0,  1,  2,  3,  4],
        [ 5,  6,  7,  8,  9],
        [10, 11, 12, 13, 14]],
      [[15, 16, 17, 18, 19],
        [20, 21, 22, 23, 24],
        [25, 26, 27, 28, 29]]])
'''

b = np.array([[True, True, False], [False, True, True]])
x[b]
'''
array([[ 0,  1,  2,  3,  4],
      [ 5,  6,  7,  8,  9],
      [20, 21, 22, 23, 24],
      [25, 26, 27, 28, 29]])
'''

高级索引和基本索引结合

当至少有一个切片（：）时，省略号（…）或者索引中的 newaxis（或者数组的维数比高级索引多），那么行为可能会更复杂。这就像连接每个高级索引元素的索引结果。

在最简单的情况下，只有一个高级索引与一个切片相结合。例如：

y = np.arange(35).reshape(5,7)
y[np.array([0, 2, 4]), 1:3]
'''
array([[ 1,  2],
       [15, 16],
       [29, 30]])
'''

实际上，切片和索引数组操作是独立的。切片操作提取索引为1和2的列（即第二列和第三列），然后是索引数组操作，该操作提取索引为0、2和4的行（即第一行、第三行和第五行）。这相当于：

y[:, 1:3][np.array([0, 2, 4]), :]
'''
array([[ 1,  2],
       [15, 16],
       [29, 30]])
'''

例如，一个高级索引可以替换一个切片，结果数组将是相同的。但是，它是一个副本，可能具有不同的内存布局。如果可能，最好是切片。例如：

x = np.array([[ 0,  1,  2],
              [ 3,  4,  5],
              [ 6,  7,  8],
              [ 9, 10, 11]])
x[1:2, 1:3]
# array([[4, 5]])
x[1:2, [1, 2]]
# array([[4, 5]])

理解多个高级索引组合的最简单方法可能是根据生成的形状进行思考。索引操作有两个部分，由基本索引（不包括整数）定义的子空间和高级索引部分定义的子空间。需要区分两种索引组合的情况：

高级索引由切片、省略号(Ellipsis)或newaxis分隔。例如 x[arr1, :, arr2]。
高级指数彼此相邻。例如 x[..., arr1, arr2, :]，但不是 x[arr1, :, 1]，因为1是这方面的高级索引。

在第一种情况下，高级索引操作产生的维度首先出现在结果数组中，然后是子空间维度。在第二种情况下，来自高级索引操作的维度被插入到结果数组中与它们在初始数组中的相同位置（后一种逻辑使简单的高级索引的行为就像切片一样）。

案例1

假设 x.shape 是 (10, 20, 30) 并且 ind 是一个 (2, 3, 4) 形的索引 intp 数组，那么 result = x[..., ind, :] 的形状是 (10, 2, 3 , 4, 30) 因为 (20,) 形子空间已被 (2, 3, 4) 形广播索引子空间替换。如果我们让 i, j, k 在 (2, 3, 4) 形子空间上循环，那么 result[..., i, j, k, :] = x[..., ind[i, j, k ]，:]。此示例产生与 x.take(ind, axis=-2) 相同的结果。

案例2

令 x.shape 为 (10, 20, 30, 40, 50) 并假设 ind_1 和 ind_2 可以广播到形状 (2, 3, 4)。那么 x[:, ind_1, ind_2] 的形状为 (10, 2, 3, 4, 40, 50)，因为来自 X 的 (20, 30) 形子空间已被替换为来自 X 的 (2, 3, 4) 子空间指数。然而，x[:, ind_1, :, ind_2] 具有形状 (2, 3, 4, 10, 30, 50)，因为在索引子空间中没有明确的放置位置，因此它被附加到开头。总是可以使用 .transpose() 将子空间移动到任何需要的地方。请注意，无法使用 take 复制此示例。

案例3

切片可以与广播的布尔索引结合使用：

x = np.arange(35).reshape(5, 7)
b = x > 20
b
'''
array([[False, False, False, False, False, False, False],
      [False, False, False, False, False, False, False],
      [False, False, False, False, False, False, False],
      [ True,  True,  True,  True,  True,  True,  True],
      [ True,  True,  True,  True,  True,  True,  True]])
'''

x[b[:, 5], 1:3]
'''
array([[22, 23],
      [29, 30]])
'''

参考

https://numpy.org/devdocs/user/basics.indexing.html

<< 基础索引 NumPy 教程 np.select() 多条件选择元素 >>