如何计算DataFrame中每个obj1除自身外的最小距离对应对象？

我有一个存储三个对象间距离的DataFrame：

import pandas as pd
df = pd.DataFrame({'obj1': ['a', 'a', 'a', 'b', 'b', 'b', 'c', 'c', 'c'], 
                   'obj2': ['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c'],
                   'dist': [0, 1, 2, 1, 0, 3, 2, 3, 0]})

需要计算每个obj1除自身外距离最小的对应对象，预期返回结果为[('a', 'b', 1), ('b', 'a', 1), ('c', 'a', 2)]。

可以通过以下步骤实现需求：

1. 过滤自身距离记录：排除obj1与obj2相同的行（即距离为0的情况），只保留不同对象间的距离数据。
2. 分组筛选最小距离：按obj1分组后，提取每组中dist值最小的行。
3. 转换为目标格式：将筛选结果转换为预期的元组列表。

具体代码如下：

import pandas as pd

# 初始化DataFrame
df = pd.DataFrame({'obj1': ['a', 'a', 'a', 'b', 'b', 'b', 'c', 'c', 'c'], 
                   'obj2': ['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c'],
                   'dist': [0, 1, 2, 1, 0, 3, 2, 3, 0]})

# 过滤掉对象自身到自身的距离行
filtered_df = df[df['obj1'] != df['obj2']]

# 按obj1分组，获取每组中dist最小的行
min_dist_rows = filtered_df.loc[filtered_df.groupby('obj1')['dist'].idxmin()]

# 转换为元组列表
result = list(min_dist_rows.itertuples(index=False, name=None))

print(result)
# 输出：[('a', 'b', 1), ('b', 'a', 1), ('c', 'a', 2)]

代码细节说明 #

• df[df['obj1'] != df['obj2']]：精准排除自身距离数据，避免后续计算干扰。
• filtered_df.groupby('obj1')['dist'].idxmin()：找到每组中距离最小值对应的行索引，确保取到的是每组距离最小的记录。
• itertuples(index=False, name=None)：直接将DataFrame行转为无额外信息的元组，快速得到目标格式。

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