Pandas多索引Series基于另一Series值筛选及分组的优化方案咨询

原始数据定义 #

import pandas as pd

s1 = pd.Series({(x,y):x*y for x in range(4) for y in range(4)})
s2 = pd.Series({('a',0):0,('a',1):2,('a',2):3,('b',0):0,('b',1):3})

需求说明 #

需要基于s2的值筛选s1的第一级索引，生成如下结构的三级索引Series，后续用于分组操作：对s2的每个条目，按s1索引的第1级分组取最大值，再按第1级分组求和。

目标Series结构：

a  0  0    0
      1    0
      2    0
      3    0
   2  0    0
      1    2
      2    4
      3    6
b  1  0    0
      1    1
      2    2
      3    3
   3  0    0
      1    3
      2    6
      3    9
dtype: int64

当前实现方式 #

目前通过字典推导式结合pd.concat实现：

s3 = pd.concat({idx:s1.loc[s2.loc[idx]] for idx in s2.index.get_level_values(0)})

优化方案 #

直接遍历s2的所有键值对，用s2的二元索引作为上层层级，筛选对应s1的结果后拼接，再调整索引层级顺序即可，代码更简洁直观：

s3 = pd.concat({idx: s1.loc[v] for idx, v in s2.items()}).reorder_levels([1, 0, 2]).sort_index()

方案解释 #

1. 遍历s2的每个键值对：idx是s2的二元索引（如('a',0)），v是对应的值（如0）
2. 用v筛选s1的第一级索引，得到对应子集
3. 通过pd.concat将所有子集按idx拼接为三级索引Series
4. 用reorder_levels调整索引层级顺序为[s2第一级索引, s2第二级索引, s1第二级索引]，最后排序索引匹配目标结构

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