Pandas如何按一二级列索引筛选多级列索引DataFrame

测试数据构造 #

先复现你给出的示例DataFrame，方便验证结果：

import pandas as pd

# 构建两级列索引
columns = pd.MultiIndex.from_arrays(
    [['A', 'A', 'B', 'B'], ['one', 'two', 'three', 'four']]
)
# 生成示例数据
df = pd.DataFrame(
    [
        [0.895717, -1.206412, 1.431256, -1.170299],
        [0.410835, 0.132003, -0.076467, 1.130127],
        [-1.413681, 1.024180, 0.875906, 0.974466]
    ],
    index=pd.Index(['A', 'B', 'C'], name='ID'),
    columns=columns
)

需求1：按一级列索引筛选该层级下所有二级列 #

如果不需要保留一级列层级，直接传入一级索引名即可，注意字符串类型的索引名需要加引号：

# 直接筛选，结果仅保留二级列索引
res1 = df['A']

如果需要和你给出的示例一样，保留原有的两级列结构，使用loc做列切片即可：

# 保留两级列结构的筛选结果
res1 = df.loc[:, ['A']]

返回结果和预期完全一致：

A          
           one       two
ID                      
A     0.895717 -1.206412
B     0.410835  0.132003
C    -1.413681  1.024180

需求2：跨一级索引，匹配不同一级索引下的指定二级列 #

Pandas原生没有df[A[two],B[three,four]]这种语法，最直接的实现方式是传入由(一级索引名, 二级索引名)格式元组组成的列表，Pandas会自动匹配对应多级列：

# 选取A下的two列、B下的three和four列
res2 = df[
    [
        ('A', 'two'),
        ('B', 'three'),
        ('B', 'four')
    ]
]

返回结果和预期完全一致：

A         B          
           two     three      four
ID                                
A    -1.206412  1.431256 -1.170299
B     0.132003 -0.076467  1.130127
C     1.024180  0.875906  0.974466

如果需要高频做这类筛选，可以写一个简单的辅助函数，实现接近你预期的传参风格：

def select_multi_cols(df, *col_rules):
    """
    多级列筛选辅助函数
    col_rules格式：(一级索引名, 二级索引名/二级索引名列表)
    """
    selected_cols = []
    for level1, level2 in col_rules:
        if isinstance(level2, list):
            selected_cols.extend([(level1, l2) for l2 in level2])
        else:
            selected_cols.append((level1, level2))
    return df[selected_cols]

# 调用示例
res2 = select_multi_cols(df, ('A', 'two'), ('B', ['three', 'four']))

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