NumPy reduceat 高效查找分段子数组最大值

本文详细介绍了如何利用 numpy 的 `np.maximum.reduceat` 函数，以完全矢量化的方式高效查找根据指定索引划分的子数组的最大值。该方法避免了显式数组分割和 python 循环，通过巧妙地构造索引数组（包含起始索引0），实现了对大型数据集的性能优化和代码简洁性。

场景描述：分段数组的最大值查找

在数据处理中，我们经常会遇到需要将一个一维 NumPy 数组按照一系列指定索引进行逻辑上的分段，并对每个分段（子数组）执行聚合操作（例如求最大值、最小值或总和）。一个直观但效率较低的方法是先使用 np.split 显式分割数组，然后通过 Python 循环遍历每个子数组并计算其最大值。

例如，给定一个数组 arr 和分割点 ind：

import numpy as np

arr = np.arange(12)
ind = np.array([3, 5, 9])

# 传统方法：先分割，再循环
sub_arrays = np.split(arr, ind)
# 结果: [array([0, 1, 2]), array([3, 4]), array([5, 6, 7, 8]), array([ 9, 10, 11])]

max_values_loop = [sub_array.max() for sub_array in sub_arrays]
# 结果: [2, 4, 8, 11]
print(f"传统循环方法获取的最大值: {max_values_loop}")

虽然这种方法易于理解，但当数组规模较大或分割点众多时，显式的数组分割和 Python 循环会引入显著的性能开销。因此，寻找一种纯 NumPy 风格的矢量化解决方案变得尤为重要。

矢量化解决方案：np.maximum.reduceat

NumPy 提供了一系列强大的矢量化操作，其中 np.ufunc.reduceat 函数是解决此类分段聚合问题的理想工具。reduceat 函数能够沿着指定轴，在由索引定义的不同“切片”上应用一个通用函数 (ufunc)。对于查找分段最大值，我们可以使用 np.maximum.reduceat。

np.maximum.reduceat 的工作原理

np.maximum.reduceat 函数接受两个主要参数：待操作的数组和一组索引。它会从每个索引位置开始，将通用函数（在这里是 maximum）应用到下一个索引位置（或数组末尾）之前的元素上。换句话说，reduceat(arr, indices) 将会计算 arr[indices[i]:indices[i+1]] 的最大值，直到最后一个索引。

关键注意事项：索引数组的构造

为了确保 np.maximum.reduceat 能够正确地处理从数组开头开始的第一个分段，我们必须在 ind 数组的开头手动添加索引 0。这是因为 reduceat 期望 indices 数组中的每个元素都指示一个新分段的起始位置。如果 0 不存在，第一个分段将无法被正确识别。

Cutout.Pro

AI驱动的视觉设计平台

331

查看详情

实施细节与示例

结合上述原理，我们可以将之前的例子完全矢量化实现如下：

# 导入必要的库
import numpy as np

# 原始数组
arr = np.arange(12)
# 分割点索引
ind = np.array([3, 5, 9])

# 构造 reduceat 所需的完整索引数组
# 将 0 添加到原始分割点数组的开头
full_indices = np.concatenate(([0], ind))
print(f"构造后的完整索引数组: {full_indices}")
# 结果: [ 0  3  5  9]

# 使用 np.maximum.reduceat 查找每个分段的最大值
max_values_vectorized = np.maximum.reduceat(arr, full_indices)
print(f"矢量化方法获取的最大值: {max_values_vectorized}")
# 结果: [ 2  4  8 11]

通过上述代码，我们可以看到 max_values_vectorized 的结果与传统循环方法 max_values_loop 完全一致，但整个过程没有显式的数组分割和 Python 循环，从而实现了更高的执行效率。

优势与应用场景

使用 np.maximum.reduceat 的矢量化方法具有以下显著优势：

• 性能提升： 对于大型数组，NumPy 的 C 实现能够显著优于 Python 循环，尤其是在数据量大时。
• 代码简洁： 将多行循环代码精简为一行，提高了代码的可读性和维护性。
• 内存效率： 避免了创建中间子数组列表，减少了内存开销。

这种方法不仅限于查找最大值。reduceat 可以与任何 NumPy 通用函数（ufunc）结合使用，例如 np.add.reduceat 用于求和，np.minimum.reduceat 用于求最小值，np.multiply.reduceat 用于求乘积等。这使得 reduceat 成为处理分段聚合任务的强大且灵活的工具，特别适用于时间序列分析、信号处理或任何需要对数据流进行窗口化或分段处理的场景。

总结

np.maximum.reduceat 提供了一种高效、矢量化的方式来计算由指定索引定义的分段子数组的最大值。其关键在于正确构造包含起始索引 0 的索引数组。掌握这一技术，能够显著提升 NumPy 数据处理的性能和代码质量，是专业数据分析和科学计算中不可或缺的技能。

以上就是NumPy reduceat 高效查找分段子数组最大值的详细内容，更多请关注php中文网其它相关文章！