跳到主要内容

NumPy 数组索引和切片

数组索引和切片是 NumPy 中最常用的操作之一。本章将详细介绍各种索引和切片方法。

基本索引​

一维数组索引​

一维数组的索引方式与 Python 列表类似：

import numpy as np

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

# 正向索引（从0开始）
print(f"arr[0]: {arr[0]}")   # 第一个元素: 0
print(f"arr[3]: {arr[3]}")   # 第四个元素: 3

# 负向索引（从-1开始）
print(f"arr[-1]: {arr[-1]}") # 最后一个元素: 9
print(f"arr[-3]: {arr[-3]}") # 倒数第三个元素: 7

# 索引越界会报错
# print(arr[10])  # IndexError

多维数组索引​

多维数组使用逗号分隔的索引元组：

import numpy as np

# 创建二维数组
arr = np.array([[1, 2, 3, 4],
                [5, 6, 7, 8],
                [9, 10, 11, 12]])

print(f"数组:\n{arr}")

# 使用逗号分隔各维度索引
print(f"arr[0, 0]: {arr[0, 0]}")   # 第一行第一列: 1
print(f"arr[1, 2]: {arr[1, 2]}")   # 第二行第三列: 7
print(f"arr[-1, -1]: {arr[-1, -1]}") # 最后一行最后一列: 12

# 只指定部分索引
print(f"arr[0]: {arr[0]}")           # 整行，返回一维数组
print(f"arr[:, 0]: {arr[:, 0]}")    # 整列，返回一维数组

# 三维数组
arr_3d = np.array([[[1, 2], [3, 4]],
                   [[5, 6], [7, 8]]])
print(f"三维数组形状: {arr_3d.shape}")  # (2, 2, 2)
print(f"arr_3d[0]:\n{arr_3d[0]}")   # 第一"页"
print(f"arr_3d[0, 1]: {arr_3d[0, 1]}")  # 第一页第二行
print(f"arr_3d[0, 1, 0]: {arr_3d[0, 1, 0]}")  # 3

切片操作​

一维数组切片​

切片语法为 start:stop:step，与 Python 列表完全一致：

import numpy as np

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

# 基本切片
print(f"arr[2:5]: {arr[2:5]}")     # [2, 3, 4] - 索引2到4（不含5）
print(f"arr[:5]: {arr[:5]}")      # [0, 1, 2, 3, 4] - 从开头到4
print(f"arr[5:]: {arr[5:]}")      # [5, 6, 7, 8, 9] - 从5到结尾
print(f"arr[:]: {arr[:]}")        # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] - 全部

# 负索引切片
print(f"arr[-3:]: {arr[-3:]}")    # [7, 8, 9] - 最后三个
print(f"arr[:-3]: {arr[:-3]}")    # [0, 1, 2, 3, 4, 5, 6] - 除了最后三个

# 步长切片
print(f"arr[::2]: {arr[::2]}")     # [0, 2, 4, 6, 8] - 每隔一个
print(f"arr[1::2]: {arr[1::2]}")   # [1, 3, 5, 7, 9] - 从1开始
print(f"arr[::-1]: {arr[::-1]}")   # [9, 8, 7, 6, 5, 4, 3, 2, 1, 0] - 反转
print(f"arr[5:2:-1]: {arr[5:2:-1]}")  # [5, 4, 3] - 反向切片

# 省略参数
print(f"arr[::]: {arr[::]}")       # 全部（等同于 arr[:]）
print(f"arr[::3]: {arr[::3]}")    # [0, 3, 6, 9] - 步长3

多维数组切片​

多维数组的切片是每个维度独立切片的组合：

import numpy as np

arr = np.array([[1, 2, 3, 4, 5],
                [6, 7, 8, 9, 10],
                [11, 12, 13, 14, 15],
                [16, 17, 18, 19, 20]])

print(f"数组 (4x5):\n{arr}")

# 各维度独立切片
print(f"arr[1:3, 2:4]:\n{arr[1:3, 2:4]}")
# 行1-2，列2-3

# 省略某些维度
print(f"arr[:, 0]: {arr[:, 0]}")    # 所有行，第0列
print(f"arr[0, :]: {arr[0, :]}")   # 第0行，所有列
print(f"arr[:, :]:\n{arr[:, :]}")  # 全部

# 步长切片
print(f"arr[::2, ::2]:\n{arr[::2, ::2]}")
# 行间隔2，列间隔2

# 负索引切片
print(f"arr[-2:, -3:]:\n{arr[-2:, -3:]}")
# 最后2行，后3列

切片是视图而非副本​

这是 NumPy 与 Python 列表的重要区别：切片返回的是视图，而不是副本：

import numpy as np

arr = np.array([1, 2, 3, 4, 5])
sliced = arr[1:4]

print(f"原数组: {arr}")
print(f"切片: {sliced}")

# 修改切片会影响原数组
sliced[0] = 99
print(f"修改切片后，原数组: {arr}")  # [1, 99, 3, 4, 5]

# 如果需要副本，使用 copy()
arr2 = np.array([1, 2, 3, 4, 5])
sliced_copy = arr2[1:4].copy()
sliced_copy[0] = 99
print(f"修改副本后，原数组: {arr2}")  # [1, 2, 3, 4, 5] - 不变

高级索引​

整数数组索引​

使用整数数组作为索引，可以自由选择任意元素：

import numpy as np

arr = np.array([10, 20, 30, 40, 50, 60])

# 一维整数数组索引
indices = [0, 2, 4]
print(f"arr[[0, 2, 4]]: {arr[indices]}")  # [10, 30, 50]

# 重复索引
indices2 = [0, 0, 1, 1, 2, 2]
print(f"重复索引: {arr[indices2]}")  # [10, 10, 20, 20, 30, 30]

# 二维索引
arr_2d = np.array([[1, 2, 3],
                    [4, 5, 6],
                    [7, 8, 9]])

# 分别指定行索引和列索引
rows = [0, 1, 2]
cols = [2, 1, 0]
print(f"arr[rows, cols]: {arr_2d[rows, cols]}")  # [3, 5, 7]

# 选择特定位置的元素
print(f"arr[[0, 0], [0, 2]]: {arr_2d[[0, 0], [0, 2]]}")  # [1, 3]

布尔索引​

使用布尔数组作为索引，选择符合条件的元素：

import numpy as np

arr = np.array([10, 20, 30, 40, 50, 60, 70, 80])

# 创建布尔条件
condition = arr > 40
print(f"条件: {condition}")  # [False False False False False True True True]

# 使用布尔数组索引
print(f"arr[condition]: {arr[condition]}")  # [50, 60, 70, 80]

# 常用写法：直接使用条件表达式
print(f"arr[arr > 40]: {arr[arr > 40]}")  # [50, 60, 70, 80]
print(f"arr[(arr > 20) & (arr < 60)]: {arr[(arr > 20) & (arr < 60)]}")  # [30, 40, 50]

# 二维数组布尔索引
arr_2d = np.array([[1, 2, 3],
                    [4, 5, 6],
                    [7, 8, 9]])

# 选择大于5的元素（返回一维数组）
print(f"arr_2d[arr_2d > 5]: {arr_2d[arr_2d > 5]}")  # [6, 7, 8, 9]

# 按行条件选择
print(f"arr_2d[arr_2d[:, 0] > 3]:\n{arr_2d[arr_2d[:, 0] > 3]}")
# 选择第一列大于3的行

np.where 函数​

np.where 可以根据条件返回值或索引：

import numpy as np

arr = np.array([10, 20, 30, 40, 50, 60])

# 基础用法：条件索引
indices = np.where(arr > 30)
print(f"大于30的索引: {indices}")  # (array([3, 4, 5]),)
print(f"大于30的元素: {arr[indices]}")  # [40, 50, 60]

# 三参数用法：类似三元表达式
result = np.where(arr > 30, arr * 2, arr)
print(f"三参数 where: {result}")  # [10, 20, 30, 80, 100, 120]

# 二维数组
arr_2d = np.array([[1, 2, 3],
                    [4, 5, 6],
                    [7, 8, 9]])

# 返回满足条件的索引
rows, cols = np.where(arr_2d > 5)
print(f"大于5的位置: 行{rows}, 列{cols}")

# 使用 where 创建掩码
mask = arr_2d > 5
print(f"掩码:\n{mask}")

花式索引​

花式索引（Fancy Indexing）使用整数数组或布尔数组进行索引：

import numpy as np

arr = np.arange(20).reshape(4, 5)
print(f"原数组:\n{arr}")

# 整数数组索引 - 选择特定行
rows = [0, 2, 3]
print(f"选择行 [0, 2, 3]:\n{arr[rows]}")

# 整数数组索引 - 选择特定列
cols = [1, 3, 4]
print(f"选择列 [1, 3, 4]:\n{arr[:, cols]}")

# 组合使用
print(f"arr[rows][cols]:\n{arr[rows][:, cols]}")

# 更高效的方式：arr[rows, cols[:, np.newaxis]]
print(f"arr[rows, cols]:\n{arr[np.ix_(rows, cols)]}")

ix_ 函数​

np.ix_ 用于创建网格索引，使得花式索引更清晰：

import numpy as np

arr = np.arange(12).reshape(3, 4)
print(f"原数组:\n{arr}")

# 选择特定行和列的交叉点
rows = [0, 2]
cols = [1, 3]
print(f"交叉点:\n{arr[np.ix_(rows, cols)]}")

# ix_ 确保生成正确的广播索引
# 等价于:
# arr[[0, 2]][:, [1, 3]]

高级索引函数​

NumPy 提供了一组专门的索引函数，用于更灵活地提取和修改数组元素。这些函数在处理复杂数据操作时非常有用。

np.take - 沿轴提取元素​

np.take 沿指定轴从数组中提取元素，比花式索引更灵活：

import numpy as np

arr = np.array([4, 3, 5, 7, 6, 8])

# 基本用法：提取指定索引位置的元素
indices = [0, 1, 4]
result = np.take(arr, indices)
print(f"提取索引 [0, 1, 4] 的元素: {result}")  # [4, 3, 6]

# 等价的花式索引写法
print(f"等价的花式索引: {arr[indices]}")

# 二维数组沿轴提取
arr_2d = np.arange(12).reshape(3, 4)
print(f"\n二维数组:\n{arr_2d}")

# 沿 axis=0 提取（提取行）
rows_to_take = [0, 2]
result_axis0 = np.take(arr_2d, rows_to_take, axis=0)
print(f"\n提取第0、2行:\n{result_axis0}")

# 沿 axis=1 提取（提取列）
cols_to_take = [1, 3]
result_axis1 = np.take(arr_2d, cols_to_take, axis=1)
print(f"提取第1、3列:\n{result_axis1}")

# 不指定 axis 时，先展平再提取
result_flat = np.take(arr_2d, [0, 5, 11])
print(f"\n展平后提取 [0, 5, 11]: {result_flat}")

mode 参数：处理越界索引

import numpy as np

arr = np.array([10, 20, 30, 40, 50])

# 'raise' 模式（默认）：越界报错
try:
    result = np.take(arr, [0, 10])
except IndexError as e:
    print(f"'raise' 模式越界报错: {e}")

# 'clip' 模式：越界索引截断到有效范围
result_clip = np.take(arr, [0, 10], mode='clip')
print(f"\n'clip' 模式: {result_clip}")  # 10 被截断为 4，取 arr[4]=50

# 'wrap' 模式：越界索引循环映射
result_wrap = np.take(arr, [0, 6, 12], mode='wrap')
print(f"'wrap' 模式: {result_wrap}")  # 6%5=1, 12%5=2

np.take_along_axis - 匹配索引提取​

np.take_along_axis 根据索引数组和数据数组的形状匹配来提取元素，常与 np.argsort 配合使用：

import numpy as np

# 场景：对每行按值排序，获取排序后的值
arr = np.array([[3, 1, 4],
                [6, 5, 2]])

# 获取每行的排序索引
sort_indices = np.argsort(arr, axis=1)
print(f"排序索引:\n{sort_indices}")

# 使用 take_along_axis 获取排序后的值
sorted_arr = np.take_along_axis(arr, sort_indices, axis=1)
print(f"\n每行排序后:\n{sorted_arr}")

# 类似地，按列排序
sort_indices_col = np.argsort(arr, axis=0)
sorted_arr_col = np.take_along_axis(arr, sort_indices_col, axis=0)
print(f"\n每列排序后:\n{sorted_arr_col}")

# 获取每行的前 k 小元素
k = 2
topk_indices = np.argpartition(arr, k, axis=1)[:, :k]
topk_values = np.take_along_axis(arr, topk_indices, axis=1)
print(f"\n每行最小的 {k} 个元素:\n{topk_values}")

np.choose - 按索引选择构造数组​

np.choose 根据索引数组从多个选项中构造新数组：

import numpy as np

# 基本用法：根据索引选择元素
choices = [[0, 1, 2, 3],
           [10, 11, 12, 13],
           [20, 21, 22, 23]]
indices = [2, 0, 1, 2]  # 选择第2、0、1、2个数组的对应位置元素

result = np.choose(indices, choices)
print(f"选择结果: {result}")  # [20, 1, 12, 23]

# 解释：result[0]=choices[2][0]=20, result[1]=choices[0][1]=1, ...

# 二维索引数组
indices_2d = [[0, 1],
              [2, 0]]
result_2d = np.choose(indices_2d, choices)
print(f"\n二维选择:\n{result_2d}")

# mode 参数处理越界索引
indices_oob = [0, 1, 2, 3, 4]  # 4 越界（只有 0-2 三个选项）
result_clip = np.choose(indices_oob, choices, mode='clip')
print(f"\n'clip' 模式处理越界: {result_clip}")  # 4 被截断为 2

result_wrap = np.choose(indices_oob, choices, mode='wrap')
print(f"'wrap' 模式处理越界: {result_wrap}")  # 4 % 3 = 1

实战示例：颜色映射

import numpy as np

# 假设有分类标签（0, 1, 2, 3）
labels = np.array([[0, 1, 2],
                   [3, 1, 0]])

# 每个类别对应的颜色
colors = [
    [255, 0, 0],    # 类别0: 红色
    [0, 255, 0],    # 类别1: 绿色
    [0, 0, 255],    # 类别2: 蓝色
    [255, 255, 0]   # 类别3: 黄色
]

# 使用 choose 将标签映射为颜色
image = np.choose(labels, colors)
print(f"标签映射为颜色:\n{image}")

np.select - 多条件选择​

np.select 根据多个条件从多个选项中选择，比嵌套 np.where 更清晰：

import numpy as np

# 成绩数据
scores = np.array([55, 72, 85, 91, 68, 45, 88])

# 定义条件和对应值
conditions = [
    scores < 60,      # 不及格
    scores < 70,      # 及格
    scores < 80,      # 中等
    scores < 90,      # 良好
    scores >= 90      # 优秀
]

choices = [
    'F',  # 不及格
    'D',  # 及格
    'C',  # 中等
    'B',  # 良好
    'A'   # 优秀
]

# 根据条件选择等级
grades = np.select(conditions, choices, default='N/A')
print(f"成绩: {scores}")
print(f"等级: {grades}")

# 使用 default 参数设置默认值
scores2 = np.array([55, 72, 85, 91])
conditions2 = [
    scores2 < 60,
    scores2 >= 80
]
choices2 = ['不及格', '优秀']
result = np.select(conditions2, choices2, default='中等')
print(f"\n带默认值: {result}")  # [不及格, 中等, 优秀, 优秀]

实战示例：数据分箱

import numpy as np

# 年龄数据
ages = np.array([5, 12, 18, 25, 35, 50, 65, 80])

# 年龄分组
conditions = [
    ages < 13,                    # 儿童
    (ages >= 13) & (ages < 20),   # 青少年
    (ages >= 20) & (ages < 40),   # 青年
    (ages >= 40) & (ages < 60),   # 中年
    ages >= 60                    # 老年
]

groups = ['儿童', '青少年', '青年', '中年', '老年']

age_groups = np.select(conditions, groups)
print("年龄分组:")
for age, group in zip(ages, age_groups):
    print(f"  {age}岁: {group}")

np.extract - 提取满足条件的元素​

np.extract 从数组中提取满足条件的元素：

import numpy as np

arr = np.arange(12).reshape(3, 4)
print(f"原数组:\n{arr}")

# 提取大于 5 的元素
condition = arr > 5
extracted = np.extract(condition, arr)
print(f"\n大于 5 的元素: {extracted}")

# 与布尔索引等价
print(f"布尔索引结果: {arr[condition]}")

# 提取偶数
even_mask = arr % 2 == 0
even_numbers = np.extract(even_mask, arr)
print(f"\n偶数元素: {even_numbers}")

np.put - 替换指定位置的元素​

np.put 原地替换数组中指定位置的元素（在展平后的数组上操作）：

import numpy as np

arr = np.arange(5)
print(f"原数组: {arr}")

# 替换索引 0 和 2 位置的元素
np.put(arr, [0, 2], [-44, -55])
print(f"替换后: {arr}")  # [-44, 1, -55, 3, 4]

# 二维数组（展平后操作）
arr_2d = np.arange(12).reshape(3, 4)
print(f"\n二维数组:\n{arr_2d}")

np.put(arr_2d, [0, 5, 11], [100, 200, 300])
print(f"替换索引 0, 5, 11 后:\n{arr_2d}")

# mode 参数处理越界索引
arr3 = np.arange(5)
np.put(arr3, 10, -1, mode='clip')  # 索引 10 被截断为 4
print(f"\n'clip' 模式: {arr3}")

arr4 = np.arange(5)
np.put(arr4, 7, -1, mode='wrap')  # 索引 7 % 5 = 2
print(f"'wrap' 模式: {arr4}")

np.place - 根据条件替换元素​

np.place 根据布尔掩码替换数组中的元素：

import numpy as np

arr = np.arange(12).reshape(3, 4)
print(f"原数组:\n{arr}")

# 将大于 5 的元素替换为 -1
np.place(arr, arr > 5, [-1])
print(f"\n大于 5 的替换为 -1:\n{arr}")

# 使用多个替换值（循环使用）
arr2 = np.arange(10)
np.place(arr2, arr2 % 2 == 0, [100, 200])
print(f"\n偶数位置替换为 100, 200 循环: {arr2}")

# 复杂条件
arr3 = np.random.randint(0, 100, 15)
print(f"\n随机数组: {arr3}")
np.place(arr3, (arr3 > 50) & (arr3 < 80), [-1, -2])
print(f"50-80之间的替换: {arr3}")

np.putmask - 根据掩码替换元素​

np.putmask 与 np.place 类似，但替换值需要与数组形状匹配：

import numpy as np

arr = np.arange(6)
print(f"原数组: {arr}")

# 创建掩码
mask = arr % 2 == 0  # 偶数位置

# 替换值数组
values = np.array([100, 200, 300])

np.putmask(arr, mask, values)
print(f"偶数位置替换: {arr}")  # [100, 1, 200, 3, 300, 5]

# putmask vs place 的区别
arr1 = np.arange(6)
arr2 = np.arange(6)

mask = arr1 > 2
values = [10, 20, 30]

# place: 只使用前 N 个值（N 是掩码中 True 的数量）
np.place(arr1, mask, values)
print(f"\nplace 结果: {arr1}")  # [0, 1, 2, 10, 20, 30]

# putmask: 值数组需要能广播到原数组形状
np.putmask(arr2, mask, np.array([10, 20, 30]))
print(f"putmask 结果: {arr2}")

np.put_along_axis - 沿轴替换元素​

np.put_along_axis 是 np.take_along_axis 的逆操作：

import numpy as np

# 创建数组
arr = np.zeros((3, 4))
print(f"原数组:\n{arr}")

# 要放置的值和索引
values = np.array([[1, 2, 3, 4],
                   [5, 6, 7, 8],
                   [9, 10, 11, 12]])

# 每行放置值的列索引
indices = np.array([[0, 1, 2, 3],
                    [3, 2, 1, 0],
                    [1, 3, 0, 2]])

# 沿 axis=1 放置值
np.put_along_axis(arr, indices, values, axis=1)
print(f"\n沿轴放置后:\n{arr}")

# 实际应用：将每行的最大值设为 0
arr2 = np.random.rand(3, 4)
print(f"\n随机数组:\n{arr2}")

# 找每行最大值的列索引
max_indices = np.argmax(arr2, axis=1, keepdims=True)
print(f"每行最大值位置: {max_indices.T}")

# 将最大值设为 0
np.put_along_axis(arr2, max_indices, 0, axis=1)
print(f"最大值设为 0 后:\n{arr2}")

高级索引函数对照表​

函数操作类型说明
np.take提取沿轴提取元素，支持越界处理
np.take_along_axis提取按匹配索引提取，常配合 argsort
np.choose构造按索引从选项中构造数组
np.select构造多条件选择，替代嵌套 where
np.extract提取提取满足条件的元素
np.put修改原地替换指定位置元素
np.place修改按条件替换元素
np.putmask修改按掩码替换元素
np.put_along_axis修改沿轴按索引放置元素

索引和切片的选择​

场景推荐方法说明
选择连续区域切片 arr[1:3, 2:4]返回视图，性能好
选择特定元素整数索引 arr[[0, 2], [1, 3]]返回副本
按条件选择布尔索引 arr[arr > 0]灵活方便
复杂选择np.where + 整数索引组合使用
import numpy as np

arr = np.arange(12).reshape(3, 4)
print(f"数组:\n{arr}")

# 选择连续区域（推荐切片）
region = arr[1:, 2:]
print(f"切片区域:\n{region}")

# 选择分散位置（使用花式索引）
disperse = arr[[0, 2], [0, 3]]
print(f"分散位置: {disperse}")

# 复杂条件（使用布尔索引）
complex_condition = arr[(arr > 3) & (arr < 10)]
print(f"复杂条件: {complex_condition}")

小结​

本章介绍了 NumPy 的各种索引和切片方法：

  1. 基本索引：使用整数索引访问单个元素
  2. 切片操作：使用 start:stop:step 语法，切片返回视图
  3. 高级索引：
    • 整数数组索引：自由选择元素，返回副本
    • 布尔索引：按条件筛选
  4. np.where：根据条件获取索引或值
  5. 花式索引：使用整数数组进行复杂选择

理解这些索引方法的区别（视图 vs 副本）对于正确操作 NumPy 数组非常重要。