时间戳

首页 > 解决方案 > 如何在 groupby 函数上应用 nsmallest

python - 如何在 groupby 函数上应用 nsmallest

问题描述

我有一个包含三列的数据框 - ID, Name, Type.

我使用以下代码使用 ID 和 Type 对数据框进行了排序 df_MI2 = df_MI1.sort_values(by=['ID', 'TYPE'])

所以，现在我有一个看起来像这样的数据框。

+--------+--------+-------+
|ID      |Name    |Type   |
|12      |A       |1      |
|12      |A       |2      |
|12      |A       |3      |
|12      |A       |4      |
|13      |B       |2      |
|13      |B       |4      |
|14      |C       |4      |
|15      |D       |3      |
|15      |D       |4      |
|15      |D       |5      |
|16      |D       |6      |
+--------+--------+-------+

现在，我只想为每个 ID 选择类型的前两个值。我怎样才能做到这一点？

我尝试了以下方法：

  1. df_MI3 = df_MI2.groupby('ID').nsmallest(2, 'TYPE')

这给了我一个错误。

2. df_MI3 = df_MI2.groupby('ID').min() 这给我"TYPE"每个人一个"ID"

标签： python

解决方案

使用cumsum，cumcount并创建一个新列作为increment。

然后删除所有行 where increment > 2。

请在下面找到示例

import pandas as pd
import numpy as np

d = {'ID' : ['I2', 'I2', 'I2', 'I2', 'I3', 'I3', 'I4', 'I5', 'I5', 'I5', 'I6'],
        'Name': ['A', 'A', 'A', 'A', 'B', 'B', 'C', 'D', 'D', 'D', 'D'],
        'Type': [1,2,3,4,2,4,4,3,4,5,6]}
df = pd.DataFrame(d)
m=df['ID']
b = m.cumsum()
df['increment'] = np.where(m, df.groupby(m.ne(m.shift()).cumsum()).cumcount()+1, 0)
df = df[df['increment'] < 3]
df.drop('increment', 1, inplace=True)
print(df)

输出是

    ID Name  Type
0   I2    A     1
1   I2    A     2
4   I3    B     2
5   I3    B     4
6   I4    C     4
7   I5    D     3
8   I5    D     4
10  I6    D     6
[Finished in 1.2s]

推荐阅读