优化基于DataFrame索引匹配的分析类型筛选Python实现方案

输入数据 #

import pandas as pd

# 用户提供的数据
fdf = pd.DataFrame(
    columns=['user_data'],
    data=[10,14,1],
    index=['alpha','beta','gamma']
)
# fdf 内容：
#         user_data
# alpha          10
# beta           14
# gamma           1

# 预设的分析支持表
sdf = pd.DataFrame(
    columns=['AD_analysis','BGD_analysis','ABG_analysis'],
    data=[[1,0,1],[0,1,1],[0,1,1],[1,1,0]],
    index=['alpha','beta','gamma','delta']
)
# sdf 内容：
#          AD_analysis  BGD_analysis  ABG_analysis
# alpha              1             0             1
# beta               0             1             1
# gamma              0             1             1
# delta              1             1             0

sdf的规则：某列对应所有用户提供的索引值均为1时，该分析即可运行。例如AD_analysis需要alpha和delta同时存在于fdf索引中才能运行。

需求 #

根据fdf的索引，筛选出sdf中可运行的分析类型，预期输出为['ABG_analysis']。

原实现问题 #

原代码依赖大量if判断，写法不够Pythonic，还存在逻辑错误（如误用位运算&代替逻辑运算and、索引存在性判断写法错误）：

# 原实现代码
com_idx = fdf.index.intersection(sdf.index)

if len(com_idx)==3 & com_idx.isin('alpha'):
    print('ABG_analysis')
if len(com_idx)==3 & com_idx.isin('delta'):
    print('BGD_analysis')
if len(com_idx)==2 :
    print('AD_analysis')

优化实现方案 #

利用Pandas向量运算特性，通过列求和直接筛选符合条件的分析类型，代码简洁且扩展性强：

import pandas as pd

# 1. 获取两个DataFrame的共同索引
com_idx = fdf.index.intersection(sdf.index)
# 2. 筛选sdf中共同索引对应的行
filtered_sdf = sdf.loc[com_idx]
# 3. 筛选出所有共同索引都支持的分析类型（列和等于共同索引数量）
valid_analyses = filtered_sdf.columns[filtered_sdf.sum() == len(com_idx)].tolist()

print(valid_analyses)  # 输出: ['ABG_analysis']

逻辑说明 #

• sdf中1表示对应索引支持该分析，当某列的和等于共同索引的数量时，说明所有用户提供的索引都支持该分析，该分析即可运行。
• 此方案无需硬编码索引数量或具体索引名称，新增分析类型或索引时无需修改代码。

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