如何避免重复调用np.searchsorted实现带偏移的数组堆叠？

我现在用np.searchsorted在数组里定位数值，想把得到的结果和它加1后的数组堆叠起来。目前我是这么写的：

import numpy as np
a = np.array([(1, 3.5), (1, 2.1), (1, 5.8), (1, 0.)])
b = np.arange(0.5, 5.5, 1.)
c = np.searchsorted(b, a[:, 1])
d = np.column_stack((c, c + 1))

我想要更直接的写法，不用重复调用np.searchsorted，比如类似np.column_stack((np.searchsorted(...), np.searchsorted(...) + 1))但只调用一次函数，预期结果是[[3 4] [2 3] [5 6] [0 1]]，有没有更好的实现思路？

几种简洁的实现思路 #

其实不用重复调用searchsorted，只要把单次调用的结果做简单的数组变形或运算就能实现，给你推荐几个实用的方法：

1. 广播运算（最简洁高效）

利用NumPy的广播特性，把一维的结果转成二维列向量后，直接和偏移量[0,1]相加，一步生成目标数组：

import numpy as np
a = np.array([(1, 3.5), (1, 2.1), (1, 5.8), (1, 0.)])
b = np.arange(0.5, 5.5, 1.)

# 只调用一次searchsorted
c = np.searchsorted(b, a[:, 1])
# 转成列向量后加偏移量，自动广播成二维数组
d = c[:, None] + [0, 1]

这里c[:, None]把形状为(4,)的一维数组变成(4,1)的二维列向量，和长度为2的数组[0,1]广播相加后，直接得到形状为(4,2)的目标数组，完全符合预期。

2. 使用海象运算符实现单行写法（Python3.8+）

如果想把代码压缩到一行，同时保证只调用一次searchsorted，可以用Python3.8引入的海象运算符:=来临时存储结果：

d = np.column_stack([(c := np.searchsorted(b, a[:, 1])), c + 1])

这个写法在一行内完成了结果存储和堆叠，可读性也不错，适合喜欢简洁代码的场景。

3. 拼接二维数组（思路直观）

如果你更习惯用拼接类的函数，也可以把原数组和偏移后的数组分别转成二维列向量，再用concatenate拼接：

c = np.searchsorted(b, a[:, 1])
d = np.concatenate([c[:, None], (c + 1)[:, None]], axis=1)

本质和广播方法类似，都是通过维度变换来实现堆叠，结果完全一致。

验证结果 #

不管用哪种方法，最终输出的d都是：

array([[3, 4],
       [2, 3],
       [5, 6],
       [0, 1]])

完全符合你的预期，而且都只调用了一次np.searchsorted，避免了不必要的重复计算。

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