在pandas列值和值列表之间高效查找

Question

我有一个包含n个元素的列表比如:

[5,30,60,180,240]

和具有以下特征的数据帧

id1 id2 feat1
1     1   40
1     2   40
1     3   40
1     4   40
2     6   87
2     7   87
2     8   87

id1 + id2的组合是唯一的，但所有具有相同id1的记录共享feat1的值。我想写一个函数通过groupby + apply(或任何更快的)来运行它，创建一个名为"close_number"的列。'closest_number'将是给定id1+id2(或id1作为记录共享的记录)的fe1列与列表中的每个元素之间最近的元素。

所需输出:

id1 id2 feat1 closest_number
1     1   40      30
1     2   40      30
1     3   40      30
1     4   40      30
2     6   87      60
2     7   87      60
2     8   87      60

如果这将是一个标准的2数组查找问题，我可以这样做:

def get_closest(array, values):
# make sure array is a numpy array
array = np.array(array)
# get insert positions
idxs = np.searchsorted(array, values, side="left")

# find indexes where previous index is closer
prev_idx_is_less = ((idxs == len(array))|(np.fabs(values - array[np.maximum(idxs-1, 0)]) < np.fabs(values - array[np.minimum(idxs, len(array)-1)])))
idxs[prev_idx_is_less] -= 1

return array[idxs]

如果我对这里的列应用这个，我将得到输出:

array([30, 60])

但是，我将无法获得关于它们与30和60对应的索引的任何信息。

做这件事的最佳方法是什么?由于我的元素列表非常小，我在我的数据集中创建了距离列，然后我选择了一个让我获得最小距离的列。

但是我认为应该有一种更优雅的方式来做到这一点。

BRE

Answer 1

使用get_closest如下:

# obtain the series with index id1 and values feat1
vals = df.groupby("id1")["feat1"].first().rename("closest_number")
# find the closest values and assign them back
vals[:] = get_closest(s, vals)
# merge the series into the original DataFrame
res = df.merge(vals, right_index=True, left_on="id1", how="left")
print(res)

id1  id2  feat1  closest_number
0    1    1     40              30
1    1    2     40              30
2    1    3     40              30
3    1    4     40              30
4    2    6     87              60
5    2    7     87              60
6    2    8     87              60