如何根据两个lambda条件筛选组并基于这些条件创建新列？

Question

这是我的数据表：

import pandas as pd

df = pd.DataFrame(
    {
        'a': ['x', 'x', 'x', 'x', 'y', 'y', 'y', 'y', 'z', 'z', 'z', 'p', 'p', 'p', 'p'],
        'b': [1, -1, 1, 1, -1, 1, 1, -1, -1, -1, -1, 1, 1, 1, 1]
    }
)

这是我期望的输出结果。我想要创建一个名为 c 的新列：

    a  b    c
0   x  1    first
1   x -1    first
2   x  1    first
3   x  1    first
4   y -1    second
5   y  1    second
6   y  1    second
7   y -1    second
11  p  1    first
12  p  1    first
13  p  1    first
14  p  1    first

分组是通过列 a 来定义的。我想要筛选 df，选择那些第一行的 b 是 1 或者第二行的 b 是 1 的组。

我用以下代码实现了这个：

df1 = df.groupby('a').filter(lambda x: (x.b.iloc[0] == 1) | (x.b.iloc[1] == 1))

然后为了给 df1 创建列 c，同样是通过 a 来定义组。如果每个组的第一行 b 是 1，那么 c 就是 first；如果第二行 b 是 1，那么 c 就是 second。

需要注意的是，对于组 p，第一行和第二行的 b 都是 1，对于这些组，我希望 c 是 first。

也许我处理这个问题的方法完全是错误的。

Answer 1

使用 GroupBy.cumcount 来进行计数，先在 b 中筛选出只有 1 的数据，然后通过字典映射只保留第一个或第二个 1，接着用 Series.dropna 删除不匹配的行，再通过 DataFrame.join 添加 a 列，最后去掉重复项，并用 DataFrame.merge 将结果添加到原始的数据框中：

s = df.groupby('a').cumcount()[df['b'].eq(1)].map({0: 'first', 1: 'second'}).dropna()

out = (df.merge(s.to_frame('c').join(df.a).drop_duplicates('a'), how='left')
         .dropna(subset=['c']))
print (out)
    a  b       c
0   x  1   first
1   x -1   first
2   x  1   first
3   x  1   first
4   y -1  second
5   y  1  second
6   y  1  second
7   y -1  second
11  p  1   first
12  p  1   first
13  p  1   first
14  p  1   first  

另一个想法：

s = (df.assign(g = df.groupby('a').cumcount())[df['b'].eq(1)]
       .drop_duplicates('a').set_index('a')['g'])

out = df.assign(c = df['a'].map(s.map({0: 'first', 1: 'second'}))).dropna(subset=['c'])
print (out)
    a  b       c
0   x  1   first
1   x -1   first
2   x  1   first
3   x  1   first
4   y -1  second
5   y  1  second
6   y  1  second
7   y -1  second
11  p  1   first
12  p  1   first
13  p  1   first
14  p  1   first

Answer 2

我觉得使用 transform 也可以在这种情况下帮上忙-

df["c"] = df.groupby("a")["b"].transform(lambda x: "first" if x.iloc[0] == 1 else ("second" if x.iloc[1] == 1 else None))
df.dropna()

输出结果 ->

    a   b   c
0   x   1   first
1   x   -1  first
2   x   1   first
3   x   1   first
4   y   -1  second
5   y   1   second
6   y   1   second
7   y   -1  second
11  p   1   first
12  p   1   first
13  p   1   first
14  p   1   first

每次循环大约需要 1.09 毫秒，误差是 ± 119 微秒（这是10次运行的平均值和标准差，每次运行10次）

如果你想把所有操作放在一行里 -

df = df.assign(c=df.groupby("a")["b"].transform(lambda x: "first" if x.iloc[0] == 1 else ("second" if x.iloc[1] == 1 else None))).dropna()

但这样会增加时间到 - 每次循环大约需要 1.24 毫秒，误差是 ± 357 微秒（同样是10次运行的平均值和标准差，每次运行10次）

Answer 3

这是一个通用的方法，可以处理任意数量的第一个 1 的位置：

d = {0: 'first', 1: 'second'}

s = (df.groupby('a')['b']
       .transform(lambda g: g.reset_index()[g.values==1]
                  .first_valid_index())
       .replace(d)
     )

out = df.assign(c=s).dropna(subset=['c'])

注意事项：

• 如果你去掉 replace 这一步，你会在 c 中得到一个整数
• 如果用 map 替代 replace，你可以忽略那些没有被定义为字典键的位置

输出结果：

    a  b       c
0   x  1   first
1   x -1   first
2   x  1   first
3   x  1   first
4   y -1  second
5   y  1  second
6   y  1  second
7   y -1  second
11  p  1   first
12  p  1   first
13  p  1   first
14  p  1   first

来自评论的示例：

df = pd.DataFrame({'a': ['x', 'x', 'x', 'x', 'y', 'y', 'y', 'y', 'z', 'z', 'z', 'p', 'p', 'p', 'p'],
                  'b': [1, -1, 1, 1, -1, 1, 1, -1, -1, -1, 1, 1, 1, 1, 1]})

d = {0: 'first', 1: 'second'}

s = (df.groupby('a')['b']
       .transform(lambda g: g.reset_index()[g.values==1]
                  .first_valid_index())
       .map(d)
     )

out = df.assign(c=s).dropna(subset=['c'])

    a  b       c
0   x  1   first
1   x -1   first
2   x  1   first
3   x  1   first
4   y -1  second
5   y  1  second
6   y  1  second
7   y -1  second
11  p  1   first
12  p  1   first
13  p  1   first
14  p  1   first

你也可以只过滤出特定的行：

m1 = df.groupby('a').cumcount().le(1)
m2 = df['b'].eq(1)
out = df.loc[df['a'].isin(df.loc[m1&m2, 'a'])]