python - 从 Pandas 数据框中的组内选择特定行

Question

假设我有一个如下所示的数据框:

   group  level
0      1     10
1      1     10
2      1     11
3      2      5
4      2      5
5      3      9
6      3      9
7      3      9
8      3      8

期望的输出是这样的:

   group  level
0      1     10
5      3      9

也就是说，逻辑是这样的:查看每个组，如果 level 列中存在多个不同的值，则返回该组中的第一行。例如，2 组中没有行被选中，因为 level 列中存在的唯一值是 5。

此外，如果我想要最后一行而不是第一行这样的组，情况会如何变化？

我尝试将 group_by 语句与从 level 列中的条目创建集合相结合，但未能产生任何几乎合理的结果。

Answer 1

这可以通过 groupby 并使用 apply 在每个组上运行一个简单的函数来完成:

def get_first_val(group):
    has_multiple_vals = len(group['level'].unique()) >= 2
    if has_multiple_vals:
        return group['level'].loc[group['level'].first_valid_index()]
    else:
        return None

df.groupby('group').apply(get_first_val).dropna()
Out[8]: 
group
1    10
3     9
dtype: float64

还有一个 last_valid_index() 方法，因此您不必进行任何巨大的更改以获取最后一行。

如果您有其他想要保留的列，您只需要稍微调整一下:

import numpy as np
df['col1'] = np.random.randint(10, 20, 9)
df['col2'] = np.random.randint(20, 30, 9)
df
Out[17]: 
   group  level  col1  col2
0      1     10    19    21
1      1     10    18    24
2      1     11    14    23
3      2      5    14    26
4      2      5    10    22
5      3      9    13    27
6      3      9    16    20
7      3      9    18    26
8      3      8    11    2

def get_first_val_keep_cols(group):
    has_multiple_vals = len(group['level'].unique()) >= 2
    if has_multiple_vals:
        return group.loc[group['level'].first_valid_index(), :]
    else:
        return None

df.groupby('group').apply(get_first_val_keep_cols).dropna()
Out[20]: 
       group  level  col1  col2
group                          
1          1     10    19    21
3          3      9    13    27