如何按已知时间值索引二维NumPy数组并提取对应质量数据

需求说明 #

我有一个左侧为时间（天数）、右侧为质量的二维NumPy数组，需要筛选出对应1到10年（按天数计算）的10组目标数据，预期输出如下：

array([[ 3.64990e+02, -4.09452e+06],
       [ 7.29730e+02, -5.45551e+06],
       [ 1.09450e+03, -6.39009e+06],
       [ 1.45910e+03, -7.15813e+06],
       [ 1.82380e+03, -7.87166e+06],
       [ 2.18860e+03, -8.58857e+06],
       [ 2.55340e+03, -9.33691e+06],
       [ 2.91820e+03, -1.01254e+07],
       [ 3.28310e+03, -1.09508e+07],
       [ 3.64790e+03, -1.18056e+07]])

原始数据示例 #

原始数组由times和mass数组堆叠而成，首尾数据如下：

n = np.stack((times, mass), axis=1)
print(n)
[[ 2.77990e-01 -3.58090e+01]
 [ 5.55970e-01 -9.10160e+01]
 [ 8.33960e-01 -1.70890e+02]
 ...
 [ 3.64960e+03 -1.18096e+07]
 [ 3.64990e+03 -1.18103e+07]
 [ 3.65000e+03 -1.18106e+07]

尝试的错误代码及问题 #

我尝试用以下代码索引目标值但未成功：

for ti, mi in n:
    for mi in [365,365*2,365*3,365*4,365*5,365*6,365*7,365*8,365*9,365*10]:
        print(ti,mi)

输出出现大量重复，且右侧质量值完全错误：

365 -35.809
730 -35.809
1095 -35.809
1460 -35.809
1825 -35.809
2190 -35.809
2555 -35.809
2920 -35.809
3285 -35.809
3650 -35.809
365 -91.016
730 -91.016
1095 -91.016
1460 -91.016
1825 -91.016
2190 -91.016
2555 -91.016
2920 -91.016
3285 -91.016
3650 -91.016
...

问题出在嵌套循环逻辑错误：外层循环遍历数组的每个元素（ti是天数，mi是质量），但内层循环重新给mi赋值为目标天数，导致每次都把当前元素的ti和所有目标天数打印，不仅重复，还覆盖了原本的质量值。

解决方案 #

利用NumPy的searchsorted函数，可以高效找到每个目标年份天数对应的最接近数据点：

import numpy as np

# 生成1-10年对应的目标天数（按365天/年计算）
target_days = np.array([365 * i for i in range(1, 11)])
# 提取数组中的时间列
time_column = n[:, 0]
# 找到每个目标天数对应的插入位置，取前一个索引（即最接近且不超过目标天数的元素）
indices = np.searchsorted(time_column, target_days, side='right') - 1
# 提取对应的数据行
result = n[indices]
print(result)

这段代码的逻辑：

1. 先生成1到10年对应的目标天数数组（365、730...3650）
2. np.searchsorted会找到每个目标天数应该插入到时间列的位置，side='right'表示找到第一个大于目标值的位置，减1后就是最后一个小于等于目标天数的元素索引
3. 通过索引提取对应的数据行，就能得到和预期一致的结果

如果你的时间列是严格递增的（从示例看是递增的），这个方法完全适用，而且效率远高于循环遍历。

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