NumPy reduceat：高效向量化查找分段数组最大值

本教程深入探讨了如何使用 NumPy 的 `np.maximum.reduceat` 函数，以高效、向量化的方式查找一维数组中逻辑分段子数组的最大值，而无需先进行显式拆分。文章将详细介绍 `reduceat` 的工作原理，并强调在索引数组中包含起始索引 `0` 的重要性，通过具体示例展示其简洁且高性能的实现。

在数据处理和科学计算中，我们经常需要对数组的特定分段（或逻辑子数组）执行聚合操作，例如查找最大值、最小值或求和。当这些分段由一系列索引定义时，一个常见的直观做法是先使用 numpy.split 将原数组拆分成子数组列表，然后遍历这些子数组进行操作。然而，对于大型数据集，这种方法涉及显式的数组拆分和 Python 循环，效率往往不高，与 NumPy 的向量化设计理念相悖。

本教程将介绍一种更高效、更符合 NumPy 风格的解决方案：利用 np.ufunc.reduceat 函数。这个强大的工具能够直接在原始数组上执行分段聚合，避免了中间的数组创建和循环开销。

理解 np.ufunc.reduceat

np.ufunc.reduceat(array, indices) 是 NumPy 通用函数（ufunc）的一个特殊方法。它允许你在 array 上，从 indices 中指定的每个位置开始，应用 ufunc（例如 maximum, sum, add 等）进行累积操作。具体来说，它会在 indices 中的每个索引处启动一个新的“规约”（reduction）过程，并将该规约的结果存储在输出数组的对应位置。

对于查找分段最大值的问题，我们将使用 np.maximum.reduceat。它的核心思想是：

• 在 indices 数组中指定的每个位置开始计算一个最大值。
• 这个最大值计算会持续到下一个 indices 中指定的索引位置之前，或者直到数组的末尾。
• 最终结果数组的长度将与 indices 数组的长度相同。

关键：准备正确的起始索引

为了让 np.maximum.reduceat 正确地处理所有分段，特别是第一个分段，我们必须确保 indices 数组包含了所有分段的起始索引。这意味着，如果你的分段是从数组的第一个元素开始的，那么 indices 数组中必须包含 0。

万相营造

阿里妈妈推出的AI电商营销工具

168

查看详情

假设我们有一个一维数组 arr 和一个表示分割点的索引数组 ind。ind 中的值表示在这些位置之后进行分割。例如，ind = [3, 5, 9] 意味着在索引 3 之前、索引 5 之前、索引 9 之前进行分割。为了使用 reduceat，我们需要提供每个分段的 起始 索引。

• 第一个分段从索引 0 开始。
• 第二个分段从 ind[0] (即 3) 开始。
• 第三个分段从 ind[1] (即 5) 开始。
• 第四个分段从 ind[2] (即 9) 开始。

因此，我们需要将原始的 ind 数组与 [0] 进行拼接，形成一个新的索引数组，作为 reduceat 的第二个参数。

示例：向量化查找分段最大值

让我们通过一个具体的例子来演示如何使用 np.maximum.reduceat。

import numpy as np

# 原始一维数组
arr = np.arange(12)
print("原始数组 arr:", arr)
# 输出: 原始数组 arr: [ 0  1  2  3  4  5  6  7  8  9 10 11]

# 分割点索引 (表示在这些索引处之后进行分割)
# 逻辑上将 arr 分割为:
# [0, 1, 2] -> max = 2
# [3, 4]    -> max = 4
# [5, 6, 7, 8] -> max = 8
# [9, 10, 11] -> max = 11
ind = np.array([3, 5, 9])
print("分割点索引 ind:", ind)
# 输出: 分割点索引 ind: [3 5 9]

# 准备用于 reduceat 的完整起始索引数组
# 必须包含 0 作为第一个分段的起始点
# 拼接后得到 [0, 3, 5, 9]
start_indices = np.concatenate(([0], ind))
print("用于 reduceat 的起始索引:", start_indices)
# 输出: 用于 reduceat 的起始索引: [0 3 5 9]

# 使用 np.maximum.reduceat 查找每个分段的最大值
max_values = np.maximum.reduceat(arr, start_indices)

print("每个分段的最大值:", max_values)
# 预期输出: 每个分段的最大值: [ 2  4  8 11]

代码解释:

1. arr = np.arange(12) 创建了一个从 0 到 11 的一维数组。
2. ind = np.array([3, 5, 9]) 定义了逻辑上的分割点。
3. start_indices = np.concatenate(([0], ind)) 是最关键的一步。它将 0（代表 arr 的起始索引）与 ind 数组连接起来，生成 [0, 3, 5, 9]。这个新的数组 start_indices 准确地告诉 reduceat 每个分段的起始位置。
   • 第一个分段从索引 0 开始，到索引 3 之前结束（即 arr[0:3]）。
   • 第二个分段从索引 3 开始，到索引 5 之前结束（即 arr[3:5]）。
   • 第三个分段从索引 5 开始，到索引 9 之前结束（即 arr[5:9]）。
   • 第四个分段从索引 9 开始，到数组末尾结束（即 arr[9:]）。
4. np.maximum.reduceat(arr, start_indices) 执行了向量化的分段最大值查找。它返回一个数组，其中每个元素对应于 start_indices 中相应索引所开始的分段的最大值。

优点与注意事项

• 性能提升: 相较于 np.split 结合 Python 列表推导式，reduceat 完全在 C 语言层面执行，避免了 Python 循环的开销和中间子数组的内存分配，对于大型数组具有显著的性能优势。
• 代码简洁性: 解决方案简洁明了，一行代码即可实现复杂的分段聚合逻辑。
• 通用性: reduceat 不仅限于 maximum。你可以将其与任何 NumPy 的通用函数（ufunc）结合使用，例如 np.add.reduceat 求和，np.minimum.reduceat 求最小值，np.multiply.reduceat 求乘积等。
• 索引数组要求: indices 数组必须是已排序的，并且其中的值必须在 [0, len(array)-1] 范围内。
• 多维数组: reduceat 也可以用于多维数组，但需要指定 axis 参数，本教程主要关注一维数组的场景。

总结

np.ufunc.reduceat 是 NumPy 中一个强大且高效的工具，用于在不显式分割数组的情况下执行分段聚合操作。通过正确构造包含所有分段起始索引的 indices 数组（特别是要包含 0），我们可以轻松实现如查找分段最大值等任务，从而编写出更高效、更符合 NumPy 风格的代码。掌握这一技巧将极大地提升你在处理大规模数值数据时的效率和代码质量。

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