
本文介绍一种高效方法,仅需一次 np.argsort() 操作,即可对数组按另一索引数组指定的顺序进行累积求和,避免重复排序开销,显著提升性能。
本文介绍一种高效方法,仅需一次 `np.argsort()` 操作,即可对数组按另一索引数组指定的顺序进行累积求和,避免重复排序开销,显著提升性能。
在 NumPy 数据处理中,常需根据一个“顺序定义数组”(如排名、优先级或时间戳索引)对另一个数值数组执行累积操作。例如,给定:
import numpy as np arr1 = np.array([1, 4, 2, 5, 6, 3, 2]) # 待累加的值 arr2 = np.array([1, 0, 2, 4, 3, 6, 5]) # 定义排序依据的键
目标是:先按 arr2 的升序确定处理顺序,对对应位置的 arr1 元素做累积和,再将结果“还原”回原始位置顺序,得到长度相同的输出数组。
传统写法常误用两次 argsort:
# ❌ 不必要且低效:两次 argsort sorting_indices = np.argsort(arr2) result = np.cumsum(arr1[sorting_indices])[np.argsort(sorting_indices)]
但关键洞察在于:sorting_indices 本身已是一个从新序号 → 原始索引的映射(即 sorting_indices[i] = j 表示排序后第 i 个元素来自原数组的第 j 位)。因此,对排序后累加结果 np.cumsum(arr1[sorting_indices]) 再按 sorting_indices 索引,本质是执行“逆置换”——将累积和值放回其对应原始位置:
# ✅ 正确且高效:仅一次 argsort sorting_indices = np.argsort(arr2) # [1, 0, 2, 4, 3, 6, 5] sorted_arr1 = arr1[sorting_indices] # [4, 1, 2, 6, 5, 2, 3] (按 arr2 升序排列) cumsum_sorted = np.cumsum(sorted_arr1) # [4, 5, 7, 13, 18, 20, 23] result = cumsum_sorted[sorting_indices] # 将每个累积值放回原 arr2 对应位置 # → [5, 4, 7, 18, 13, 23, 20]
✅ 输出完全匹配预期:array([ 5, 4, 7, 18, 13, 23, 20])。
为什么 cumsum_sorted[sorting_indices] 是正确的?
因为 sorting_indices[i] 给出的是 arr2[i] 在排序序列中的“来源位置”。例如 arr2[0] == 1,而 1 在 arr2 中排第 1 位(0-indexed),故 sorting_indices[0] == 1;那么 cumsum_sorted[1] 就是前两个最小键对应 arr1 元素的和(即 arr1[1] + arr1[0] = 4 + 1 = 5),这正是 result[0] 应有的值。
注意事项:
- 此方法要求
arr2中元素互异;若存在重复值,np.argsort的稳定性(默认为 ‘quicksort’)可能影响结果一致性,建议显式指定kind='stable'以确保可重现性; - 时间复杂度为 O(n log n)(主导于一次排序),空间复杂度 O(n),优于两次排序的隐式开销;
- 该模式可推广至其他顺序敏感操作(如
np.maximum.accumulate、分组累积等),核心思想是“排序→变换→逆索引”。
综上,通过理解 argsort 返回索引的语义并合理复用,可简洁、高效地完成顺序累积任务,无需冗余计算。










