如何基于化学键等价规则合并整理Pandas DataFrame列？

问题场景 #

给定包含化学键计数的DataFrame，其中方向相反的化学键（如OH与HO、CaO与OCa）为等价键，需要将这些等价键的计数合并求和，得到简化后的DataFrame。

原始DataFrame示例：

import pandas as pd

df = pd.DataFrame({
    'OH': [2, 0, 1],
    'HO': [3, 2, 2],
    'CaO': [2, 3, 3],
    'OCa': [0, 4, 0],
    'OO': [1, 5, 0],
    'NaMg': [1, 2, 0],
    'MgNa': [1, 0, 0]
})

输出如下：

OH  HO  CaO  OCa  OO  NaMg  MgNa
0   2   3    2    0   1     1     1
1   0   2    3    4   5     2     0
2   1   2    3    0   0     0     0

目标是合并等价键后得到：

OH  CaO  OO  NaMg
0   5    2   1     2
1   2    7   5     2
2   3    3   0     0

核心难点 #

• 列数超过3000，无法手动分组；
• 元素符号格式不固定（单大写字母、双字母首大写后小写），需正确拆分化学键字符串。

解决方案代码 #

import pandas as pd
import re

def split_bond(bond):
    # 正则匹配元素符号：大写字母开头，可选后续小写字母
    return re.findall(r'[A-Z][a-z]?', bond)

def get_canonical_bond(bond):
    # 生成化学键的标准标识：拆分元素后按字母排序拼接
    elements = split_bond(bond)
    return ''.join(sorted(elements))

# 1. 为每个列名生成标准键标识
canonical_map = {col: get_canonical_bond(col) for col in df.columns}

# 2. 建立标准标识到原始列名的映射，取每个组的第一个列名作为最终列名
reverse_map = {}
for col, canon in canonical_map.items():
    reverse_map.setdefault(canon, []).append(col)
final_col_names = {canon: cols[0] for canon, cols in reverse_map.items()}

# 3. 按标准标识分组求和
df_merged = df.groupby(canonical_map, axis=1).sum()

# 4. 重命名列并恢复原始列的出现顺序
df_merged = df_merged.rename(columns=final_col_names)
ordered_cols = []
seen = set()
for col in df.columns:
    final_col = final_col_names[canonical_map[col]]
    if final_col not in seen:
        ordered_cols.append(final_col)
        seen.add(final_col)
df_merged = df_merged[ordered_cols]

# 输出结果
print(df_merged)

代码说明 #

• 元素拆分：使用正则表达式[A-Z][a-z]?可以精准匹配所有元素符号，解决单/双字母符号的拆分问题；
• 标准标识生成：通过对拆分后的元素按字母排序拼接，确保方向相反的化学键得到相同的标识，实现自动分组；
• 分组求和：利用pandas的groupby功能批量合并等价列，高效处理数千列的场景；
• 列名与顺序保留：保留原始DataFrame中每个等价键组的首个列名，并恢复列的原始出现顺序，符合化学命名习惯。

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