使用Python Pandas按相似性为实习生分组并分配小队编号

需求说明 #

我有一个包含25列、1000余行的实习生虚拟数据集，需要根据时区、偏好语言、兴趣小组等特征的相似性，将实习生划分为每队10人的小队，并为每个小队分配唯一的小队编号。

数据集列名（中文翻译）：

["姓名","小队编号","偏好语言","兴趣小组","年龄","城市","国家","地区","时区",
 "职业","学历","过往修读课程","实习经历","产品管理","数字营销",
 "市场调研","数字插画","产品设计","产品开发","增长营销",
 "带队经验","实习资讯","产品营销 cohort","产品设计 cohort",
 "产品开发 cohort","产品增长 cohort","每周可投入时长"]

实现步骤与代码 #

1. 导入工具库并加载数据 #

用pandas加载数据集，sklearn处理特征编码和相似度计算：

import pandas as pd
from sklearn.preprocessing import OneHotEncoder
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

# 加载数据集（替换为你的文件路径）
df = pd.read_csv("interns_dataset.csv")

2. 选择相似性判断的核心特征 #

选取需求指定的特征，补充每周可投入时长这类影响协作的特征：

# 选择用于计算相似性的特征列
similarity_features = ["时区", "偏好语言", "兴趣小组", "每周可投入时长"]
feature_df = df[similarity_features].copy()

3. 编码分类特征 #

将分类特征转为数值型，方便计算相似度：

# 分离分类特征和数值特征
cat_features = ["时区", "偏好语言", "兴趣小组"]
num_features = ["每周可投入时长"]

# 对分类特征做独热编码
encoder = OneHotEncoder(sparse_output=False, drop="first")
encoded_cats = encoder.fit_transform(feature_df[cat_features])
encoded_cat_df = pd.DataFrame(encoded_cats, columns=encoder.get_feature_names_out(cat_features))

# 合并编码后的分类特征与数值特征
processed_features = pd.concat([encoded_cat_df, feature_df[num_features].reset_index(drop=True)], axis=1)

4. 计算相似度并分组 #

提供两种实用的分组方式：

方式一：KMeans聚类分组（适合大规模数据）

先估算聚类数量，再调整分组大小确保每组10人：

from sklearn.cluster import KMeans

# 估算聚类数：总人数 / 每组人数
n_clusters = len(df) // 10

# 训练KMeans模型
kmeans = KMeans(n_clusters=n_clusters, random_state=42)
df["小队编号"] = kmeans.fit_predict(processed_features)

# 调整分组：平衡每组人数
for squad in df["小队编号"].unique():
    squad_size = len(df[df["小队编号"] == squad])
    if squad_size > 10:
        # 拆分多余人员到新组
        extra_members = df[df["小队编号"] == squad].sample(squad_size - 10)
        new_squad_num = df["小队编号"].max() + 1
        df.loc[extra_members.index, "小队编号"] = new_squad_num
    elif squad_size < 10:
        # 从人数最多的组抽取人员补充
        largest_squad = df["小队编号"].value_counts().idxmax()
        if len(df[df["小队编号"] == largest_squad]) > 10:
            transfer_members = df[df["小队编号"] == largest_squad].sample(10 - squad_size)
            df.loc[transfer_members.index, "小队编号"] = squad

方式二：贪心式相似性分组（精准匹配相似样本）

从第一个未分配样本开始，每次选择最相似的9个样本组成小队：

# 计算余弦相似度矩阵
similarity_matrix = cosine_similarity(processed_features)
# 标记已分配的样本
assigned = np.zeros(len(df), dtype=bool)
squad_num = 0

while not assigned.all():
    # 找到第一个未分配样本
    start_idx = np.argwhere(~assigned)[0][0]
    assigned[start_idx] = True
    # 获取与该样本相似度最高的前9个未分配样本
    similar_indices = np.argsort(similarity_matrix[start_idx])[::-1]
    similar_indices = [idx for idx in similar_indices if not assigned[idx]][:9]
    # 标记样本并分配小队编号
    assigned[similar_indices] = True
    df.loc[[start_idx] + similar_indices, "小队编号"] = squad_num
    squad_num += 1

5. 验证分组结果 #

检查每个小队的人数分布：

print(df["小队编号"].value_counts().describe())

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