首页 > 解决方案 > 使用 apply 函数同时填充 2 列数据框

问题描述

我有一些像这样(简化)的代码。实际的数据列表有数万大小,而不仅仅是 3 个。

有一本员工词典,我从中制作了一个 DataFrame。有一个包含附加员工信息的字典对象列表。

还:

  1. 员工列表和额外的员工信息 (master_info_list) 重叠,但每个都有各自独有的项目。
  2. 我使用的“索引”(StaffNumber)实际上在额外的人员信息中以“SN_”为前缀,因此无法直接进行比较。
  3. master_info_list 中的 StaffNumber 重复是有意的(这就是我收到它的方式!)。

我想要做的是将两个新列填充到数据框中,从额外的员工信息中获取数据。我可以通过对 get_department_and_manager 进行 2 次单独调用来做到这一点,一次用于部门,一次用于经理。这样可行。但是,它“感觉”就像我应该能够从 get_department_and_manager 的输出中获取 2 个字段并一次性填充数据帧,但我正在努力使语法正确。什么是正确的语法(如果可能)?此外,以我的方式(使用 for 循环)遍历列表似乎效率低下。有没有更好的办法?

我看到的所有示例似乎都是从数据框中的现有数据创建新列,或者它们是在比较两个“列表”(或列表和字典)之前不需要数据混搭的简单示例。

import pandas as pd


def get_department_and_manager(row, master_list):
    dept = 'bbb'
    manager = 'aaa'
    for i in master_list:
        if i['StaffNumber'] == 'SN_' + row['StaffNumber']:
            dept = i['data']['Department']
            manager = i['data']['Manager']
            break
    return [dept, manager]


staff = {'Name': ['Alice', 'Bob', 'Dave'],
         'StaffNumber': ['001', '002', '004']}

master_info_list = [{'StaffNumber': 'SN_001', 'data': {'StaffNumber': 'SN_001', 'Department': 'Sales', 'Manager': 'Luke' }},
      {'StaffNumber': 'SN_002', 'data': {'StaffNumber': 'SN_002', 'Department': 'Marketing', 'Manager': 'Mary' }},
      {'StaffNumber': 'SN_003', 'data': {'StaffNumber': 'SN_003', 'Department': 'IT', 'Manager': 'Neal' }}]

df = pd.DataFrame(data=staff)
df[['Department']['Manager']] = df.apply(get_department_and_manager, axis='columns', args=[master_info_list])

print(df)

标签: pythonpandasdataframe

解决方案


如果我理解正确,您可以使用.merge

x = pd.DataFrame([v["data"] for v in master_info_list])
x["StaffNumber"] = x["StaffNumber"].str.split("_").str[-1]
print(df.merge(x, on="StaffNumber", how="left"))

印刷:

    Name StaffNumber Department Manager
0  Alice         001      Sales    Luke
1    Bob         002  Marketing    Mary
2   Dave         004        NaN     NaN

推荐阅读