根据DataFrame中A列的值筛选保留指定数据行

我有如下结构的DataFrame：

A    B     C    D     E    F
A3   25    26   A1    36    39
A3   10    10   A4    10    10
A5   66   22    A2    26    26
A6   66   26    A3    35    36
A7   34   36    A5    66    22

希望根据A列的值筛选保留部分数据，得到如下输出（注意：相同A列值对应的B列与C列之和等于E列与F列之和）：

A    B     C    D     E    F
A3   25    26   A3    35    36
A3   10    10   A5    66    22
A5   66   22    NA    NA    NA
A6   66   26    NA    NA    NA
A7   34   36    NA    NA    NA

核心逻辑 #

先提取原始数据中D列值存在于A列的行，将这些行的(D,E,F)作为匹配资源，分配给A列对应值的行；其余行的D/E/F设为缺失值（NA）。

代码实现 #

import pandas as pd
import numpy as np

# 构建初始DataFrame
data = {
    'A': ['A3', 'A3', 'A5', 'A6', 'A7'],
    'B': [25, 10, 66, 66, 34],
    'C': [26, 10, 22, 26, 36],
    'D': ['A1', 'A4', 'A2', 'A3', 'A5'],
    'E': [36, 10, 26, 35, 66],
    'F': [39, 10, 26, 36, 22]
}
df = pd.DataFrame(data)

# 提取可用的匹配项：筛选D列值在A列中的行，转为字典列表
a_unique = set(df['A'])
match_records = df[df['D'].isin(a_unique)][['D', 'E', 'F']].to_dict('records')

# 处理需要分配匹配项的行
target_rows = df[df['A'].isin([item['D'] for item in match_records])].copy()
for idx in range(len(target_rows)):
    target_rows.iloc[idx, 3:] = pd.Series(match_records[idx])

# 处理无需分配的行，设置D/E/F为缺失值
rest_rows = df[~df['A'].isin([item['D'] for item in match_records])].copy()
rest_rows[['D', 'E', 'F']] = np.nan

# 合并结果并按A列排序
final_df = pd.concat([target_rows, rest_rows]).sort_values('A').reset_index(drop=True)
print(final_df)

运行结果 #

A   B   C    D     E     F
0  A3  25  26  A3  35.0  36.0
1  A3  10  10  A5  66.0  22.0
2  A5  66  22  NaN   NaN   NaN
3  A6  66  26  NaN   NaN   NaN
4  A7  34  36  NaN   NaN   NaN

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