如何在Pandas中返回每行距离值的最大值及其对应集群名称？

场景说明 #

已有包含user、values、group三列的Pandas DataFrame，其中values存储多值列表；已计算出各集群的质心（平均值），以及每个用户到所有集群质心的欧氏距离。需要将每行距离的最大值及其对应的集群名称整合到DataFrame中，得到指定格式的结果。

原始数据与质心计算 #

首先是数据构造与质心计算的代码：

import pandas as pd
import numpy as np
from scipy.spatial import distance

# 构造原始DataFrame
df = pd.DataFrame({'user': ['user_1', 'user_2', 'user_3', 'user_4', 'user_5',  'user_6'],
                   'values': [[1, 0, 2, 0], [1, 8, 0, 2],[6, 2, 0, 0], [5, 0, 2, 2], [3, 8, 0, 0],[6, 0, 0, 2]],
                   'group': ['B', 'A', 'C', 'A', 'B', 'B']})

# 计算各集群质心
df1 = (df.groupby('group', as_index=False)['values']
         .agg(lambda x: np.vstack(x).mean(0).round(2))
       )

整合最大距离与对应集群 #

接下来是核心实现步骤，将距离计算、最大值提取与原DataFrame合并：

# 生成每个用户到各集群质心的距离DataFrame，列名为集群名称
distance_df = pd.DataFrame(
    [
        [distance.euclidean(val, centroid) for centroid in df1['values']]
        for val in df['values']
    ],
    columns=df1['group'].tolist()
)

# 提取每行的最大距离值和对应的集群名
max_distance = distance_df.max(axis=1)
max_group = distance_df.idxmax(axis=1)

# 合并到原DataFrame，保留需要的列
result = df[['user']].assign(max_value=max_distance, group=max_group)
print(result)

输出结果 #

运行上述代码后，将得到预期的输出：

user  max_value group
0  user_1   5.744563     C
1  user_2   8.062258     C
2  user_3   4.242641     A
3  user_4   4.582576     A
4  user_5   6.708204     C
5  user_6   5.099020     A

代码说明 #

• distance_df：将每个用户的距离列表转换为以集群名为列的DataFrame，方便后续行列操作
• max(axis=1)：计算每行（每个用户）的最大距离值
• idxmax(axis=1)：获取每行最大值对应的列名（即距离最远的集群名称）
• assign()：将新生成的max_value和group列添加到仅保留user列的原DataFrame中，得到最终结果

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