MCMC高频代码优化：以精度换性能改进np.searchsorted数组匹配

问题描述 #

我有一段嵌套在MCMC内部的代码，需要执行数百万次，因此必须尽可能高效。当前代码利用已排序的arr0，通过np.searchsorted()为arr2[0]的每个元素寻找arr0中绝对最近的元素索引，然后据此向数组添加值。我希望以牺牲部分精度（容忍±0.2的误差，寻找“接近”元素而非绝对最近）换取性能提升，请问是否可行且能提升代码性能？

原始代码 #

import numpy as np
# Random initial data with the actual shapes used by my code.
Nmax = 1000000
arr0 = np.linspace(5., 30., Nmax)
D = np.random.randint(2, 4)
arr1 = np.random.uniform(-3., 3., (D, Nmax))
arr2 = np.random.uniform(10., 25., (10, 1500))

# Can these two lines be made faster?
# Indexes of elements in 'arr0' closest to the elements in 'arr2[0]'
closest_idxs = np.searchsorted(arr0, arr2[0])
# Add elements from 'arr1' to the first dimensions of 'arr2', according
# to the indexes found above.
arr_final = arr2[:arr1.shape[0]] + arr1[:, closest_idxs]

回答 #

当然可行，而且这种精度换性能的思路在你的场景下能带来非常显著的性能提升——毕竟你的代码要跑数百万次，哪怕每次循环节省几毫秒，累计下来都是巨大的时间节省。

核心优化思路：利用arr0的均匀特性跳过二分查找

你的arr0是用np.linspace生成的均匀间隔数组，这是关键突破口！np.searchsorted本质是二分查找，时间复杂度是O(M log N)（M是arr2[0]的长度，N是arr0的长度），但我们可以通过纯算术计算直接得到满足精度要求的索引，时间复杂度降到O(M)，速度快几个数量级。

具体来说：

1. 先计算arr0的间隔step和起始值offset
2. 对arr2[0]的每个元素，直接通过算术公式计算近似索引
3. 做简单的边界裁剪避免索引越界

你的arr0间隔是(30-5)/1e6 = 2.5e-5，远小于你允许的±0.2误差，所以这种近似方法得到的索引对应的arr0元素，和目标值的误差绝对满足你的要求。

优化后的代码

import numpy as np
# Random initial data with the actual shapes used by my code.
Nmax = 1000000
arr0 = np.linspace(5., 30., Nmax)
D = np.random.randint(2, 4)
arr1 = np.random.uniform(-3., 3., (D, Nmax))
arr2 = np.random.uniform(10., 25., (10, 1500))

# 提前计算arr0的间隔和起始值（可以移到MCMC循环外，只算一次）
step = arr0[1] - arr0[0]
offset = arr0[0]

# 直接计算近似索引，误差远小于±0.2
closest_idxs = ((arr2[0] - offset) / step).astype(np.int64)
# 裁剪索引到合法范围，避免越界
closest_idxs = np.clip(closest_idxs, 0, Nmax - 1)

# 后续计算和原代码一致
arr_final = arr2[:arr1.shape[0]] + arr1[:, closest_idxs]

额外优化建议

如果你的MCMC循环里arr0是固定不变的，一定要把step和offset的计算移到循环外面！只计算一次，不要每次循环都重复计算，这能再省一点时间。

非均匀arr0的拓展方案（如果后续需求变化）

如果以后arr0变成非均匀间隔的，但你依然允许±0.2的误差，可以提前对arr0按0.2的间隔做分箱：

• 预先生成每个0.2区间对应的arr0索引（比如每个区间存第一个符合条件的索引）
• 对arr2[0]的元素，直接计算它属于哪个区间，然后取出对应的索引
  这种方法同样能跳过二分查找，提升性能。

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