python - Pandas - 保持至少有两个不同代码的组

Question

我正在使用具有以下结构的 DataFrame:

import pandas as pd

df = pd.DataFrame({'group' : [1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 4, 4], 
                   'brand' : ['A', 'B', 'X', 'A', 'B', 'C', 'X', 'B', 'C', 'X', 'A', 'B'], 
                   'code' : [2185, 2185, 0, 1410, 1390, 1390, 0, 3670, 4870, 0, 2000, 0]})

print(df)

    group brand  code
0       1     A  2185
1       1     B  2185
2       1     X     0
3       2     A  1410
4       2     B  1390
5       2     C  1390
6       2     X     0
7       3     B  3670
8       3     C  4870
9       3     X     0
10      4     A  2000
11      4     B     0

我的目标是仅查看至少具有两个不同代码的组。过滤标准中不应考虑用 0 标记的缺失代码。例如，即使第 4 组中的两条记录具有不同的代码，我们也不会将这一组保留在最终的 DataFrame 中，因为其中一个代码丢失。

上面示例中生成的 DataFrame 应如下所示:

    group brand  code
1       2     A  1410
2       2     B  1390
3       2     C  1390
4       2     X     0
5       3     B  3670
6       3     C  4870
7       3     X     0

我没能对这个问题做太多事情。我认为第一步应该是创建一个掩码来删除缺少 (0) 代码的记录。像这样的东西:

mask = df['code'].eq(0)
df = df[~mask]

print(df)

    group brand  code
0       1     A  2185
1       1     B  2185
3       2     A  1410
4       2     B  1390
5       2     C  1390
7       3     B  3670
8       3     C  4870
10      4     A  2000

现在只保留至少有两个不同代码的组，但我不知道如何在 Python 中解决这个问题。另外，此方法将删除我不想要的最终 DataFrame 中缺少代码的记录。我想了解整个小组的情况。

任何额外的帮助将不胜感激。

Answer 1

这是transform():

mask = (df.groupby('group')['code']
          .transform(lambda x: x.mask(x==0)    # mask out the 0 values
                                .nunique()     # count the nunique
                    )
          .gt(1)
       )
df[mask]

输出:

   group brand  code
3      2     A  1410
4      2     B  1390
5      2     C  1390
6      2     X     0
7      3     B  3670
8      3     C  4870
9      3     X     0

选项 2:类似的想法，但没有 lambda 函数:

mask = (df['code'].mask(df['code']==0)    # mask out the 0 values
            .groupby(df['group'])         # groupby
            .transform('nunique')         # count uniques
            .gt(1)                        # at least 2
       )