如何在多层索引Pandas DataFrame中按指定上下文生成新列？

我有一个多层索引的DataFrame，构造代码如下：

import numpy as np
import pandas as pd

np.random.seed(1)
df = pd.DataFrame(
    {
        'trial':  [0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,2,2,2,2,2,2,2,2],
        't':      [0,0,0,0,1,1,1,1,2,2,2,2,3,3,3,3,0,0,0,0,1,1,1,1,2,2,2,2,3,3,3,3,0,0,0,0,1,1,1,1],
        'context':[0,1,2,3,0,1,2,3,0,1,2,3,0,1,2,3,0,1,2,3,0,1,2,3,0,1,2,3,0,1,2,3,0,1,2,3,0,1,2,3],
        'x' :np.random.rand(40),
        'y' : np.random.rand(40),
        'z' : np.random.rand(40),
        'inferred_context':   [0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,1,1,1,1,1,1,1,1],
    }
)

df = df.set_index(["trial", "t","context"])
df= df.unstack('context')
df.columns.set_names(['vals','context'],inplace=True)
df.swaplevel(0,1)

该DataFrame的行索引为两级(trial, t)，列索引为两级(vals, context)。我想要创建一个名为x_inferred_context的列，其中每行的值为对应trial和t下，inferred_context字段指定的context对应的x值。例如：

• 在(trial=0, t=3)行中，inferred_context=0，对应(vals=x, context=0)的x值为0.204452，则该行的x_inferred_context列值应为0.204452；
• 在(trial=1, t=0)行中，inferred_context=2，对应context=2的x值为0.140387，则该行的x_inferred_context列值应为0.140387。

我目前通过为每个可能的上下文创建掩码，再与df['x']相乘后求和实现了需求：

nc = np.unique((df.columns.get_level_values(level=1))).size
mask = pd.DataFrame( data = [(df['inferred_context'] == c).iloc[:,c] for c in range(nc)]).T
df['x_inferred_context'] = (mask*df['x']).sum(axis=1)

作为Pandas新手，我想询问实现该需求的标准方法是什么？是否有更简洁、可读性更强且更符合Pandas风格的实现方式？

你可以使用Pandas内置的lookup方法，这是专门为“按行匹配指定列标签取值”的场景设计的，代码简洁且可读性极强：

df['x_inferred_context'] = df['x'].lookup(df.index, df['inferred_context'])

方法说明： #

• df['x']：提取所有与x相关的列，这是一个二级列索引的DataFrame，列的二级标签对应context的取值；
• lookup方法接收两个参数：行标签（直接使用df.index，即(trial, t)的行索引）、每行需要匹配的列标签（df['inferred_context']的每行值，对应目标context）；
• 该方法会自动为每行找到对应列标签的x值，直接生成目标列。

此外，也可以用take方法实现相同效果，性能同样出色：

df['x_inferred_context'] = df['x'].take(df['inferred_context'].values, axis=1)

这两种方法都比手动构造掩码求和的写法更简洁，完全符合Pandas的风格，省去了手动生成掩码的繁琐步骤。

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