在列名内循环

Posted 2023-03-12

【中文标题】在列名内循环

【英文标题】：Loop within column name

【发布时间】：2019-03-28 07:32:17

【问题描述】：

我有一个数据框，其中包含 car_x 和 car1_y、van2_x 和 van2_y 以及 bus3_x 和 bus3_y 列。我需要一列 car1_x * car1_y + van2_x * van2_y + bus3_x * bus3_y

以下代码不起作用：

modes = 'car', 'van', 'bus'

for mode in modes:
    df['var'] = df['var_x']*df['var_y']

然后我将在 df['car']、df['van'] 和 df['bus'] 之间求和，但上面的语法是关闭的。

【参考方案1】：

要修复你的代码，你需要使用 f-strings 来让 python 知道应该插入 var 作为它的值，而不是字符串 "var"。

for mode in modes:
    df[f'var'] = df[f'var_x'] * df[f'var_y']

但这需要额外的求和步骤才能获得“结果”。

df['result'] = df[list(modes)].sum(axis=1)

---

让我们删掉额外的步骤，并在此处使用einsum 更快地执行此操作。过滤掉您的 _x 和 _y 列，然后使用 einsum 指定求和运算。

x = df.filter(like='_x')
y = df.filter(like='_y')

df['result'] = np.einsum('ij,ij->i', x, y)

感谢filter 步骤，不再需要维护单独的modes 列表。

【讨论】：

非常感谢，我错过了 f。不过你的代码更干净。

【参考方案2】：

我将使用groupby

df.groupby(df.columns.str.split('_').str[0],axis=1).prod()[['car', 'van', 'bus']].sum(1)