如何提取Multi-index DataFrame子级匹配NumPy数组的行

问题场景 #

给定如下两层索引的DataFrame：

a   b   c
0   0   0   42  65
    1   6   0   340
    2   5   71  800
    3   2   51  409
    4   0   23  279
    5   8   38  549
1   0   1   23  252
    1   9   13  977
    2   1   19  943
2   0   2   23  295
    1   3   39  458
    2   1   62  308
    3   0   95  954
    4   9   78  535
3   0   4   67  849
    1   3   46  761
    2   7   49  485
    3   0   44  638

给定NumPy数组a = np.array([2, 2, 4, 3])，需要提取每个一级索引分组下，二级索引值等于数组对应位置元素的行，期望输出为：

a   b   c
0   5   71  800
1   1   19  943
2   9   78  535
3   0   44  638

原有代码问题 #

之前尝试的代码：

i,j = df.index.levels
ix = a
df1 = pd.DataFrame(df.to_numpy()[ix])

逻辑错误点：直接用数组值作为全局行的整数位置索引取值，完全忽略多层索引的层级对应关系。不同一级分组下的二级索引是独立编号的，全局行号和分组内的子索引值没有对应关系，最终会取到错位的行。

正确实现方法 #

利用pandas多层索引的.loc标签索引能力，先构造(一级索引值, 目标二级索引值)的元组列表，直接匹配取值即可：

# 按一级索引的分组顺序，和目标数组打包成索引元组
target_index = list(zip(df.index.get_level_values(0).unique(), a))
# 标签索引提取目标行
df_result = df.loc[target_index]
# 如果不需要保留二级索引，直接删除该层级即可得到期望格式
df_result = df_result.droplevel(1)

运行后得到的结果和预期完全一致：

a   b    c
0  5  71  800
1  1  19  943
2  9  78  535
3  0  44  638

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