Python3.x 从列表的元素和python中的多个多索引数据帧执行方程

发布于03月02日

我有以下列表和数据框:

import pandas as pd
cols = pd.MultiIndex.from_tuples([("K", "A1", "A"), 
                                  ("K", "B1", "B"), 
                                  ("K", "C1", "C"),
                                  ("M", "A1", "A"), 
                                  ("M", "B1", "B"), 
                                  ("M", "C1", "C")])

data = [[5, 10, 40, 4, 8, 9], [2, 15, 70, 1, 7, 3], [6, 14, 60, 12, 41, 61]]
df = pd.DataFrame(data, columns=cols)
lst = [5, '*', 'A1.A', '+', 'C1.C']

列多索引名称在列表中列出的方式可以更改，如果可以提供帮助的话.

我想创建一个代码来执行第一个类似5 x df['A1','A'] + df['C1','C']的等式，但对于匹配第二级和第三级的所有列组合，返回以下结果:

df[("K", "N1", "N")] = 5 * df[("K", "A1", "A")] + df[("K", "C1", "C")]
df[("M", "N1", "N")] = 5 * df[("M", "A1", "A")] + df[("M", "C1", "C")]

这有可能吗？

推荐答案

您可以在第一个列级唯一键上使用pandas.eval和循环:

df_name = 'df'

for key in df.columns.get_level_values(0).unique():
    query = ''.join([f'{df_name}.{key}.{x}' if (isinstance(x, str) and '.' in x)
                     else str(x)
                     for x in lst])
    df[(key, 'N1', 'N')] = pd.eval(query)

df = df.sort_index(axis=1, level=0, sort_remaining=False)

print(df)

当然，如果您不需要从列表构建它，您可以简化为:

for key in df.columns.get_level_values(0).unique():
     df[(key, 'N1', 'N')] = 5 * df[(key, "A1", "A")] + df[(key, "C1", "C")]

输出:

   K               M             
  A1  B1  C1  N1  A1  B1  C1   N1
   A   B   C   N   A   B   C    N
0  5  10  40  65   4   8   9   29
1  2  15  70  80   1   7   3    8
2  6  14  60  90  12  41  61  121