如何在NumPy中为单调递增数组高效生成条件求和新列？

既然题目里明确说了数组的第一列是单调递增的，那我们完全可以利用这个特性来避免暴力遍历，用「二分查找+前缀和」的组合实现O(N logN)的高效解法，比O(N²)的暴力方法在大数组上快得多。

核心思路 #

1. 前缀和数组：先预处理第二列的前缀和，这样任意区间的和都可以用O(1)的时间计算出来。
2. 二分查找边界：对于每一行的阈值（当前行第一列值+10），利用单调递增的特性，用np.searchsorted快速找到第一列中第一个大于等于该阈值的位置，这个位置就是我们求和的右边界（不包含）。
3. 计算每行的和：通过前缀和数组，直接计算从当前行到右边界前一行的第二列元素之和，作为第三列的值。

代码实现 #

import numpy as np

# 示例输入数组
arr = np.array([[ 1, 5],
                [ 3, 5],
                [ 7, 5],
                [11, 5],
                [25, 5]])

# 提取第一列和第二列
col1 = arr[:, 0]
col2 = arr[:, 1]

# 构造前缀和数组，prefix[0] = 0，prefix[i]是前i个元素的和（col2[0]到col2[i-1]）
prefix = np.concatenate([[0], np.cumsum(col2)])

# 计算每行的阈值：当前行第一列值 + 10
thresholds = col1 + 10

# 用searchsorted找到每个阈值对应的右边界（第一个>=阈值的索引）
# side='left'表示返回第一个满足条件的位置
indices = np.searchsorted(col1, thresholds, side='left')

# 计算每行的求和结果：prefix[indices[n]] - prefix[n] 就是col2[n]到col2[indices[n]-1]的和
sums = prefix[indices] - prefix[np.arange(len(arr))]

# 拼接原数组和求和结果，得到最终输出
result = np.column_stack([arr, sums])

print(result)

输出验证 #

运行上面的代码，会得到和题目期望完全一致的结果：

[[ 1  5 15]
 [ 3  5 15]
 [ 7  5 10]
 [11  5  5]
 [25  5  5]]

为什么高效？ #

• np.cumsum是O(N)的时间复杂度，用于生成前缀和。
• np.searchsorted对每个元素的查找是O(logN)，整个数组处理下来是O(N logN)。
• 最后的求和计算是O(N)，所以整体时间复杂度是O(N logN)。相比暴力遍历每个行然后检查后续所有行的O(N²)，在数组规模较大时（比如10^4以上的行），性能提升会非常明显。

这个解法完全利用了题目给出的第一列单调递增的特性，是最简洁高效的NumPy实现方式之一。

内容的提问来源于stack exchange，提问作者pi_175