Skip to main content
 首页 » 编程设计

python之两个数组的高效匹配(KDTree的使用方法)

2025年05月04日9duanxz

我有两个二维数组,obs1 和 obs2。它们代表两个独立的测量系列,并且都有 dim0 = 2,并且 dim1 略有不同,比如 obs1.shape = (2, 250000) , 和 obs2.shape = (2, 250050)。 obs1[0]和obs2[0]表示时间,obs1[1]和obs2[1]表示一些空间坐标。两个数组(或多或少)都按时间排序。两个测量系列之间的时间和坐标应该相同,但实际上并非如此。此外,并非 obs1 中的每个测量值在 obs2 中都有相应的值,反之亦然。另一个问题是时间可能会有轻微的偏移。

我正在寻找一种有效的算法来将 obs2 中的最佳匹配值关联到 obs1 中的每个测量值。目前,我是这样做的:

define dt = some_maximum_time_difference 
define dx = 3 
j = 0 
i = 0 
matchresults = np.empty(obs1.shape[1]) 
for j in obs1.shape[1]: 
    while obs1[0, j] - obs2[0, j] < dt: 
        i += 1 
    matchresults[j] = i - dx + argmin(abs(obs1[1, i] - obs2[1, i-dx:i+dx+1])) 

这会产生很好的结果。但是,它非常慢,在循环中运行。

我将非常感谢有关如何提高此算法速度的想法,例如使用 KDtree 或类似的东西。

请您参考如下方法:

在这种情况下使用 cKDTree 看起来像:

from scipy.spatial import cKDTree 
 
obs2 = array with shape (2, m) 
obs1 = array with shape (2, n) 
 
kdt = cKDTree(obs2.T) 
dist, indices = kdt.query(obs1.T) 

其中 indices 将包含 obs2 中的列索引,对应于 obs1 中的每个观察值。请注意,我必须转置 obs1 和 obs2。