NumPy矢量化：高效查找子数组的最大值

本教程详细阐述如何利用numpy库的`np.maximum.reduceat`函数，以纯矢量化方式高效地计算一维数组按指定索引“分割”后各子数组的最大值。通过巧妙地构造索引数组，该方法避免了显式数组分割和python循环，显著提升了处理效率和代码的简洁性，是处理此类问题的numpy风格最佳实践。

在数据处理中，我们经常需要对一个一维NumPy数组进行逻辑上的“分割”，并对每个分割后的子数组执行聚合操作，例如查找最大值。传统的做法是先使用numpy.split将原始数组分割成子数组列表，然后遍历这个列表，对每个子数组调用其.max()方法。

考虑以下示例，我们有一个数组arr和一组分割点ind：

import numpy as np

arr = np.arange(12) # arr 为 array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11])
ind = np.array([3, 5, 9]) # 在索引3, 5, 9处进行分割

# 传统方法：先分割，再遍历查找最大值
sub_arrays = np.split(arr, ind)
# sub_arrays 结果: [array([0, 1, 2]), array([3, 4]), array([5, 6, 7, 8]), array([ 9, 10, 11])]
max_values_traditional = [sub_array.max() for sub_array in sub_arrays]
# max_values_traditional 结果: [2, 4, 8, 11]
print(f"传统方法得到的最大值: {max_values_traditional}")

尽管上述方法直观易懂，但它涉及了显式的数组分割操作和Python级别的循环。在大规模数据处理时，这种方法可能导致性能瓶颈，并且不符合NumPy倡导的矢量化计算哲学。

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

NumPy提供了一个更高效、更符合矢量化思想的解决方案：np.maximum.reduceat。这个函数是通用函数（ufunc）reduceat方法的一个特例，专门用于查找指定索引区间内的最大值。

np.ufunc.reduceat(array, indices) 的核心思想是，它会在indices数组中指定的每个起始索引处“重置”累积操作。对于np.maximum.reduceat，这意味着它会从每个指定索引开始，计算直到下一个指定索引（或数组末尾）的最大值。

Natural Language Playlist

探索语言和音乐之间丰富而复杂的关系，并使用 Transformer 语言模型构建播放列表。

67

查看详情

关键步骤：为了获得与np.split后.max()完全一致的结果，传递给reduceat的索引数组必须包含原始数组的起始索引 0。

import numpy as np

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

# 构造包含起始索引0的完整索引数组
# np.concatenate(([0], ind)) 会得到 [0, 3, 5, 9]
# 这表示我们希望计算从索引0开始到索引2的最大值，从索引3开始到索引4的最大值，
# 从索引5开始到索引8的最大值，以及从索引9开始到数组末尾的最大值。
full_indices = np.concatenate(([0], ind))

# 使用 np.maximum.reduceat 进行矢量化计算
max_values_vectorized = np.maximum.reduceat(arr, full_indices)

print(f"矢量化方法得到的最大值: {max_values_vectorized}")
# 预期输出: [ 2  4  8 11]

代码解释：

1. arr = np.arange(12): 创建一个包含0到11的NumPy数组。
2. ind = np.array([3, 5, 9]): 定义分割点。这些索引是每个子数组的起始索引（除了第一个子数组的起始索引0）。
3. full_indices = np.concatenate(([0], ind)): 这是最关键的一步。np.maximum.reduceat需要所有子数组的起始索引。由于ind只包含了后续子数组的起始索引，我们需要手动添加第一个子数组的起始索引 0。
   • [0]：代表第一个子数组的起始索引。
   • ind：代表后续子数组的起始索引。
   • np.concatenate将它们合并成 [0, 3, 5, 9]。
4. max_values_vectorized = np.maximum.reduceat(arr, full_indices):
   • reduceat会从 full_indices 中的每个索引开始，应用 maximum 操作。
   • 从索引 0 开始，计算 arr[0:3] (即 [0, 1, 2]) 的最大值，结果是 2。
   • 从索引 3 开始，计算 arr[3:5] (即 [3, 4]) 的最大值，结果是 4。
   • 从索引 5 开始，计算 arr[5:9] (即 [5, 6, 7, 8]) 的最大值，结果是 8。
   • 从索引 9 开始，计算 arr[9:] (即 [9, 10, 11]) 的最大值，结果是 11。
   • 最终结果 [2, 4, 8, 11] 正是所有子数组的最大值。

优势与注意事项

• 性能提升: np.maximum.reduceat 是在C语言层面实现的，避免了Python循环的开销，对于大型数组而言，其性能远超np.split结合列表推导式的方法。
• 内存效率: 这种方法不需要创建中间的子数组列表，从而减少了内存开销。
• 通用性: reduceat不仅限于maximum，还可以与NumPy的许多其他通用函数（如np.add.reduceat、np.sum.reduceat、np.amin.reduceat等）结合使用，实现各种分段聚合操作。
• 索引数组要求: 传递给reduceat的indices数组必须是升序的。如果索引不按升序排列，结果将是未定义的。
• 边界条件: 确保full_indices中的最后一个索引小于arr的长度。如果最后一个索引等于或大于arr的长度，reduceat会正确处理，但可能导致空切片或不符合预期的结果。

总结

当需要对NumPy数组进行逻辑分割并对每个分段执行聚合操作（如查找最大值、求和等）时，np.ufunc.reduceat 提供了一种强大且高效的矢量化解决方案。通过正确构造包含起始索引0的索引数组，我们可以避免显式分割和Python循环，从而编写出更简洁、性能更优的NumPy代码。掌握reduceat的使用，是提升NumPy编程效率和解决复杂数据处理问题的关键技能之一。

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