Python Pandas条件分组：如何将化合物按分子量约束分为每组100个

化合物分组约束实现方案 #

实现思路 #

• 先按molecular_mass升序排序全量数据，让分子量相近的化合物相邻
• 生成多个无冲突候选池，每个候选池内任意两个化合物的分子量差值均大于1，保证同池元素可共存于同一分组
• 将所有候选池的元素按顺序拼接后，每100个切分为一组，自动满足约束要求

实现代码 #

import pandas as pd

# 读取原始数据
df = pd.read_csv('data.csv')
# 按分子量升序排序
df_sorted = df.sort_values('molecular_mass', ascending=True).reset_index(drop=True)

# 构建无冲突候选池
pools = []
for _, row in df_sorted.iterrows():
    current_mass = row['molecular_mass']
    placed = False
    # 尝试放入已有候选池，仅需和池内最后一个元素比较（已排序，前面的元素更小）
    for pool in pools:
        if current_mass - pool[-1]['molecular_mass'] > 1:
            pool.append(row)
            placed = True
            break
    # 无符合要求的候选池则新建
    if not placed:
        pools.append([row])

# 可选校验：如果候选池数量超过最大分组数（1250/100≈13），说明数据无法满足约束
if len(pools) > 13:
    raise ValueError("当前数据中存在过多分子量差值在±1范围内的化合物，无法完成符合要求的分组")

# 拼接所有候选池元素并分组
flat_data = []
for pool in pools:
    flat_data.extend(pool)
df_final = pd.DataFrame(flat_data).reset_index(drop=True)

GROUP_SIZE = 100
df_final['group'] = df_final.index // GROUP_SIZE
groups = [df_final[df_final['group'] == group_id] for group_id in range(df_final['group'].max() + 1)]

方案说明 #

排序后仅需和候选池最后一个元素比较分子量差，即可保证整个候选池内无冲突，时间复杂度低。最终拼接生成的分组中，任意两个元素都来自不同的候选池，天然满足分子量差值不在±1范围内的要求。

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