python - 将参数传递给 pandas.DataFrame.transform 中的函数

Question

我编写了一个示例，其中我想将 mask_size 作为 pandas.DataFrame.transform 使用的 mask_first 中的参数传递。

import pandas as pd
import random
a = [a for a in [1,2,3,4,5] for i in range(25)]
b = [random.randrange(1, 10) for _ in range(0, 125)]
d = {'col1': a, 'col2': b}
df = pd.DataFrame(d)
def mask_first(df, mask_size):
    result = np.ones_like(df)
    result[0:mask_size] = 0
    return result
def apply_mask(df, mask_size = 10):
    mask = df.groupby(['col1'])['col1'].transform(mask_first).astype(bool)
    df = df.loc[mask]
    return df
df = apply_mask(df, mask_size = 10)

它给了我错误代码

TypeError: mask_first() missing 1 required positional argument: 'mask_size'

Answer 1

根据 pandas.DataFrame.transform 的文档，您应该能够使用 transform 方法本身将参数传递给函数:

DataFrame.transform(func, axis=0, *args, **kwargs)

哪里(摘自 pandas 文档):

*args: Positional arguments to pass to func.

**kwargs: Keyword arguments to pass to func.

尝试执行此操作将掩码大小参数传递给您的 mask_first 函数:

# Important part: .transform(mask_first, mask_size)
mask = df.groupby(['col1'])['col1'].transform(mask_first, mask_size).astype(bool)