NumPy详解：数组操作与性能优化 – wiki大全

作者： /

NumPy详解：数组操作与性能优化

NumPy (Numerical Python) 是Python中用于科学计算的核心库。它提供了一个高性能的多维数组对象（ndarray），以及用于处理这些数组的工具。NumPy的强大之处在于其底层用C和Fortran实现，使得数组操作的速度远超Python原生的列表，尤其是在处理大规模数据时。本文将详细探讨NumPy的数组操作及其性能优化策略。

一、 NumPy Ndarray 基础

1. 什么是 Ndarray？

NumPy的核心是ndarray（n-dimensional array）对象，这是一个同质的多维数组，存储着相同类型的数据项。与Python列表不同，ndarray中的所有元素必须是相同的数据类型，这使得它在内存中可以更紧凑地存储，并允许对整个数组进行批量操作。

2. 创建 Ndarray

创建ndarray的方式多种多样：

  • 从Python列表或元组创建：
    “`python
    import numpy as np
    list_data = [1, 2, 3, 4, 5]
    array_from_list = np.array(list_data)
    print(array_from_list) # 输出: [1 2 3 4 5]

    tuple_data = ((1, 2, 3), (4, 5, 6))
    array_from_tuple = np.array(tuple_data)
    print(array_from_tuple)

    输出:

    [[1 2 3]

    [4 5 6]]

    “`

  • 使用内置函数创建特定数组：

    • np.zeros(shape): 全零数组
    • np.ones(shape): 全一数组
    • np.empty(shape): 元素值随机（取决于内存状态）的数组
    • np.arange(start, stop, step): 类似range的数组
    • np.linspace(start, stop, num): 在指定区间内均匀间隔的数字
    • np.eye(N) 或 np.identity(N): 单位矩阵

    “`python
    print(np.zeros((2, 3)))

    输出:

    [[0. 0. 0.]

    [0. 0. 0.]]

    print(np.arange(0, 10, 2)) # 输出: [0 2 4 6 8]
    print(np.linspace(0, 1, 5)) # 输出: [0. 0.25 0.5 0.75 1. ]
    “`

3. 数组属性

ndarray拥有许多重要的属性：

  • .ndim: 数组的维度数量
  • .shape: 数组的维度大小（元组），例如 (rows, columns)
  • .size: 数组中元素的总数量
  • .dtype: 数组中元素的数据类型
  • .itemsize: 数组中每个元素占用的字节数
  • .data: 存储数组元素的缓冲区（通常不直接使用）

python
arr = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int32)
print(f"维度: {arr.ndim}") # 输出: 维度: 2
print(f"形状: {arr.shape}") # 输出: 形状: (2, 3)
print(f"大小: {arr.size}") # 输出: 大小: 6
print(f"数据类型: {arr.dtype}") # 输出: 数据类型: int32

二、 NumPy 数组操作

NumPy提供了丰富的数组操作功能，包括基本算术、索引、切片、形状变换、合并、拆分等。

1. 基本算术操作

NumPy数组之间的算术操作是元素级的（element-wise），这意味着操作会应用于数组中的每一个对应元素。

“`python
a = np.array([1, 2, 3])
b = np.array([4, 5, 6])

print(a + b) # 输出: [5 7 9]
print(a * b) # 输出: [ 4 10 18]
print(a * 2) # 输出: [2 4 6] (广播机制)
print(a > 2) # 输出: [False False True] (布尔数组)
矩阵乘法使用 `@` 运算符或 `np.dot()` 函数。python
matrix_a = np.array([[1, 2], [3, 4]])
matrix_b = np.array([[5, 6], [7, 8]])
print(matrix_a @ matrix_b)

输出:

[[19 22]

[43 50]]

“`

2. 索引和切片

NumPy的索引和切片与Python列表类似，但更强大，支持多维索引和布尔索引。

  • 基本索引：
    python
    arr = np.array([10, 20, 30, 40, 50])
    print(arr[0]) # 输出: 10
    print(arr[-1]) # 输出: 50

  • 切片： start:stop:step
    python
    print(arr[1:4]) # 输出: [20 30 40]
    print(arr[::2]) # 输出: [10 30 50]

  • 多维数组索引：
    python
    matrix = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
    print(matrix[0, 1]) # 输出: 2 (第一行第二列)
    print(matrix[1, :]) # 输出: [4 5 6] (第二行所有元素)
    print(matrix[:, 2]) # 输出: [3 6 9] (第三列所有元素)
    print(matrix[0:2, 1:3]) # 提取子矩阵
    # 输出:
    # [[2 3]
    # [5 6]]

  • 布尔索引： 使用布尔数组选择元素
    python
    arr = np.array([10, 20, 30, 40, 50])
    mask = arr > 25
    print(mask) # 输出: [False False True True True]
    print(arr[mask]) # 输出: [30 40 50]
    print(arr[arr % 2 == 0]) # 筛选偶数

  • 花式索引（Fancy Indexing）： 使用整数数组作为索引
    “`python
    arr = np.array([10, 20, 30, 40, 50])
    indices = np.array([0, 2, 4])
    print(arr[indices]) # 输出: [10 30 50]

    matrix = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
    rows = np.array([0, 2])
    cols = np.array([1, 0])
    print(matrix[rows, cols]) # 输出: [2 7] (提取 (0,1) 和 (2,0) 处的元素)
    “`

3. 形状变换

  • .reshape(shape): 不改变数据，只改变数组的形状。如果新的形状与原始数组的元素总数不符，会报错。
  • .ravel() 或 .flatten(): 将多维数组展平为一维数组。ravel()返回视图（可能修改原数组），flatten()返回副本。
  • .transpose() 或 .T: 数组转置。
  • np.newaxis: 增加一个维度。

“`python
arr = np.arange(12) # [ 0 1 2 3 4 5 6 7 8 9 10 11]
reshaped_arr = arr.reshape((3, 4))
print(reshaped_arr)

输出:

[[ 0 1 2 3]

[ 4 5 6 7]

[ 8 9 10 11]]

print(reshaped_arr.T) # 转置

输出:

[[ 0 4 8]

[ 1 5 9]

[ 2 6 10]

[ 3 7 11]]

增加维度

vec = np.array([1, 2, 3])
print(vec.shape) # (3,)
row_vec = vec[np.newaxis, :] # 变为行向量
print(row_vec.shape) # (1, 3)
col_vec = vec[:, np.newaxis] # 变为列向量
print(col_vec.shape) # (3, 1)
“`

4. 合并与拆分

  • 合并： np.concatenate(), np.vstack(), np.hstack(), np.dstack()

    • np.concatenate((a1, a2, ...), axis=0/1): 沿指定轴连接数组。
    • np.vstack((a1, a2)) (np.row_stack): 垂直堆叠（行方向）。
    • np.hstack((a1, a2)) (np.column_stack): 水平堆叠（列方向）。

    “`python
    a = np.array([[1, 2], [3, 4]])
    b = np.array([[5, 6], [7, 8]])

    print(np.vstack((a, b)))

    输出:

    [[1 2]

    [3 4]

    [5 6]

    [7 8]]

    print(np.hstack((a, b)))

    输出:

    [[1 2 5 6]

    [3 4 7 8]]

    “`

  • 拆分： np.split(), np.vsplit(), np.hsplit()

    • np.split(array, indices_or_sections, axis=0/1): 沿指定轴将数组拆分为子数组。

5. 通用函数 (Universal Functions – UFuncs)

UFuncs是NumPy提供的一类函数，它们对ndarray进行元素级的操作，并且速度非常快。常见的数学函数（np.sin, np.cos, np.exp, np.log, np.sqrt等）以及前面提到的算术运算符都是ufuncs。
它们可以接受任意形状的数组作为输入，并返回一个形状相同的数组作为输出。

6. 广播 (Broadcasting)

广播是NumPy中处理不同形状数组之间算术运算的一套规则。当操作两个数组时，NumPy会尝试通过复制较小数组的元素来“广播”它，使其与较大数组的形状兼容，从而执行元素级操作。

广播规则：
1. 如果两个数组的维度数量不同，则维度较小的数组的形状会在其前面补1，直到两个数组的维度数量相同。
2. 对于从后往前比较的每个维度：
* 如果两个维度大小相等，则保持不变。
* 如果其中一个维度大小为1，则将其沿着另一个维度的大小进行广播。
* 如果两个维度大小都不相等且都不为1，则会引发错误。

“`python
a = np.array([[1, 2, 3], [4, 5, 6]]) # 形状 (2, 3)
b = np.array([10, 20, 30]) # 形状 (3,)

b 会被广播成 [[10, 20, 30], [10, 20, 30]]

print(a + b)

输出:

[[11 22 33]

[14 25 36]]

c = 5 # 标量

c 会被广播成 [[5, 5, 5], [5, 5, 5]]

print(a * c)

输出:

[[ 5 10 15]

[20 25 30]]

“`
广播极大地简化了代码，避免了显式的循环，并提升了性能。

三、 性能优化

NumPy的性能优势主要来源于其底层的C实现和对向量化操作的支持。理解并利用这些特性是优化NumPy代码的关键。

1. 向量化操作 (Vectorization)

这是NumPy性能优化的核心和基石。 向量化是指将操作应用于整个数组，而不是通过Python循环逐个元素进行操作。NumPy的内部实现会自动利用高效的C/Fortran代码来执行这些批量操作，从而显著提高速度。

避免使用Python循环：
“`python

示例：计算两个数组的元素积

arr1 = np.random.rand(1000000)
arr2 = np.random.rand(1000000)

非向量化（使用Python循环） – 极慢

import time
start_time = time.time()
result_loop = [arr1[i] * arr2[i] for i in range(len(arr1))]
end_time = time.time()
print(f”Python循环耗时: {end_time – start_time:.6f} 秒”)

向量化（使用NumPy操作） – 极快

start_time = time.time()
result_vectorized = arr1 * arr2
end_time = time.time()
print(f”NumPy向量化耗时: {end_time – start_time:.6f} 秒”)
“`
在百万级元素的数组上，向量化操作的速度通常比Python循环快几个数量级。

2. 利用 UFuncs

UFuncs是NumPy中已经高度优化的向量化函数。无论何时，只要可能，就应该优先使用它们。
例如，计算数组元素的平方根，np.sqrt(arr) 远比 [math.sqrt(x) for x in arr] 更高效。

3. 选择合适的数据类型 (Dtype)

NumPy允许你为数组元素指定精确的数据类型，例如np.int8, np.int32, np.float32, np.float64等。
* 节省内存： 如果你知道数据范围较小，使用更小的数据类型（如np.int8而非np.int64）可以显著减少内存占用。
* 提高性能： 较小的数据类型可以使得CPU缓存更有效，并且在某些处理器上，对小数据类型的操作可能更快。
* 精度权衡： 浮点数类型（float32 vs float64）在精度和性能之间进行权衡。除非需要双精度，否则float32通常足够。

python
arr_large = np.arange(1000000, dtype=np.int64)
arr_small = np.arange(1000000, dtype=np.int8) # 会溢出，此处仅为示例
print(f"int64数组内存占用: {arr_large.nbytes} 字节")
print(f"int8数组内存占用: {arr_small.nbytes} 字节")

4. 避免创建不必要的中间数组 (In-place Operations / Views)

一些NumPy操作会返回原始数组的视图（view），这意味着它们共享相同的底层数据缓冲区，修改视图会影响原始数组。而有些操作则会返回副本（copy），创建新的内存空间。
* 视图操作： 切片、reshape (如果元素总数不变)、transpose 通常返回视图。
* 副本操作： 花式索引、flatten()、大部分数学运算（如arr1 + arr2）会创建新的数组。

尽量利用返回视图的操作可以减少内存分配和数据复制，从而提升性能。如果确实需要修改数组而不影响原始数组，则需要显式地进行复制，例如 arr.copy()。

5. 缓存局部性

NumPy数组在内存中是连续存储的，这有利于CPU缓存。当访问数组元素时，CPU会一次性加载一片内存到缓存中。如果你的访问模式是线性的（例如逐行或逐列遍历），那么缓存命中率会很高，性能会更好。
* 行优先与列优先： C语言和Python（默认）是行优先存储，而Fortran是列优先。在NumPy中，默认创建的数组是C-contiguous（行主序）。这意味着沿着行方向的遍历通常比沿着列方向的遍历更快。
“`python
matrix = np.random.rand(1000, 1000)

# 行遍历（更快）
start_time = time.time()
_ = np.sum(matrix, axis=1) # 沿着列求和 (对行进行操作)
end_time = time.time()
print(f"行遍历耗时: {end_time - start_time:.6f} 秒")

# 列遍历（较慢）
start_time = time.time()
_ = np.sum(matrix, axis=0) # 沿着行求和 (对列进行操作)
end_time = time.time()
print(f"列遍历耗时: {end_time - start_time:.6f} 秒")
```
(注意：这里的`np.sum(axis=...)`本身是高度优化的，所以差异可能不明显。对于更复杂的逐元素操作，这种差异会更显著。)

6. 使用 np.einsum 进行复杂张量操作

对于涉及多个数组的复杂求和、乘积、转置等操作，np.einsum提供了一个简洁而高效的接口，可以避免中间数组的创建，并允许NumPy在底层进行高度优化。它通过爱因斯坦求和约定来表达多维数组操作。

“`python
A = np.random.rand(3, 4, 5)
B = np.random.rand(4, 3, 2)

矩阵乘法，并对某些轴求和 (复杂例子)

假设我们想计算 (A_ijk * B_jil) 并对 j 求和，得到一个 (k, i, l) 的结果

result = np.einsum(‘ijk,jil->kil’, A, B)
print(result.shape) # (5, 3, 2)
``
虽然初学可能有些复杂，但掌握einsum`可以显著提升特定复杂计算的性能。

四、 总结

NumPy是Python科学计算的基石，其ndarray对象提供了高效的数据存储和操作能力。理解并熟练运用NumPy的数组操作，特别是其强大的索引、切片和广播机制，是编写高效Python数值代码的关键。

更重要的是，通过贯彻向量化原则，优先使用UFuncs，合理选择数据类型，并注意内存布局，可以最大化NumPy的性能潜力，从而在数据分析、机器学习和科学研究等领域中处理大规模数据集时获得显著的优势。避免Python循环，拥抱NumPy的向量化特性，将是您优化代码的第一步，也是最重要的一步。

滚动至顶部