时间戳

首页 > 解决方案 > Pandas：对于排序序列 B 的所有元素，查找排序序列 A 中最近元素的索引

python - Pandas：对于排序序列 B 的所有元素，查找排序序列 A 中最近元素的索引

问题描述

我有一个带有两个排序的整数列的数据框。

      A        B
0     17       15
1     18       18
2     19       20
3     20       21
4     22       21
5     23       27

对于B的所有元素，我想找到A的最接近匹配元素的索引：

      A        B       closest_match_idx
0     17       15      0
1     18       18      1
2     19       20      3
3     20       21      3
4     22       21      3
5     23       27      5

我知道我能做到

df['closest_match_idx'] = df.B.map(lambda x: (df.A - x).abs().idxmin()))

但对于显然是 O(N) 的问题，这是一个 O(N**2) 的解决方案。除了滚动我自己的索引查找功能之外，我找不到任何更好的解决方案，但这感觉就像一个有现有解决方案的问题。想法？

对于上下文，我最终要做的是为 B 的每个元素在 A 中找到最接近的匹配元素，直到最大绝对差（否则只使用 B 中的值）：

match_diff = df.B - df.A.iloc[df['closest_match_idx']]
df['output'] = B
replace_idxs = np.where(diff.abs() <= 2)
df['output'].iloc[replace_idxs] = df['A'].iloc[replace_idxs]

searchsorted 或这​​个索引技巧几乎就在那里，但并不完全。

      A        B       closest_match_idx   match_diff    output
0     17       15      0                   2             17
1     18       18      1                   0             18
2     19       20      3                   1             20
3     20       21      3                   1             20
4     22       21      3                   1             20
5     23       27      5                   4             23

我还应该注意，在我的示例中，A 和 B 具有相同的长度，但我想要一个可以推广到一系列不同长度的解决方案。

标签： pythonpandasalgorithm

解决方案

您可以使用merge_asof. 这需要对帧进行排序。这样做的好处是它支持一个tolerance参数，允许您指定一个您认为匹配的卡尺。

我将留在附加'A_match'列中，但如果您不需要它，可以将其删除。

res = pd.merge_asof(df.sort_values('B'), 
                    df.rename_axis(index='closest_idx').reset_index().drop(columns='B').sort_values('A'),
                    left_on='B', right_on='A',
                    direction='nearest',
                    suffixes=['', '_match'])

print(res)

    A   B  closest_idx  A_match
0  17  15            0       17
1  18  18            1       18
2  19  20            3       20
3  20  21            3       20
4  22  21            3       20
5  23  27            5       23

设置|距离|的容差 <= 1

res = pd.merge_asof(df.sort_values('B'), 
                    df.rename_axis(index='closest_idx').reset_index().drop(columns='B').sort_values('A'),
                    left_on='B', right_on='A',
                    direction='nearest',
                    suffixes=['', '_match'],
                    tolerance=1)

#    A   B  closest_idx  A_match
#0  17  15          NaN      NaN
#1  18  18          1.0     18.0
#2  19  20          3.0     20.0
#3  20  21          3.0     20.0
#4  22  21          3.0     20.0
#5  23  27          NaN      NaN

推荐阅读