如何從 Python 數據框中的每個類別中獲取前 n 條記錄？

Question

數據在以下數據框中的“id”列上按降序排序 -

id   Name     version     copies   price
6    MSFT       10.0        5       100   
6    TSLA       10.0        10      200
6    ORCL       10.0        15      300

5    MSFT       10.0        20      400
5    TSLA       10.0        25      500
5    ORCL       10.0        30      600

4    MSFT       10.0        35      700
4    TSLA       10.0        40      800
4    ORCL       10.0        45      900

3    MSFT       5.0         50      1000 
3    TSLA       5.0         55      1100
3    ORCL       5.0         60      1200

2    MSFT       5.0         65      1300
2    TSLA       5.0         70      1400
2    ORCL       5.0         75      1500

1    MSFT       15.0        80      1600
1    TSLA       15.0        85      1700
1    ORCL       15.0        90      1800
...

基於輸入'n'，我想過濾上面的數據，如果輸入是'2'，結果數據框應該看起來像 -

Name     version     copies   price
MSFT       10.0        5       100   
TSLA       10.0        10      200
ORCL       10.0        15      300

MSFT       10.0        20      400
TSLA       10.0        25      500
ORCL       10.0        30      600

MSFT       5.0         50      1000 
TSLA       5.0         55      1100
ORCL       5.0         60      1200

MSFT       5.0         65      1300
TSLA       5.0         70      1400
ORCL       5.0         75      1500

MSFT       15.0        80      1600
TSLA       15.0        85      1700
ORCL       15.0        90      1800

基本上，只有特定版本的“id”的前“n”組應該出現在結果數據框中。 如果一個版本的 id < n（例如，在 15.0 版中只有一個組的 id = 1），那么所有組的 id 都應該存在。

我嘗試使用groupy和head ，但它對我不起作用。 我絕對沒有其他線索可以讓這個工作。

我非常感謝您對此的任何幫助，謝謝。

Answer 1

您可以在列版本上使用groupby.transform ，並將列 id factorize為每個組的每個 id 的增量值（從 0 到 ...），然后與您的 n 進行比較並使用帶有此掩碼的loc來選擇想要的行。

n = 2
print(df.loc[df.groupby('version')['id'].transform(lambda x: pd.factorize(x)[0])<n])
    id  Name  version  copies  price
0    6  MSFT     10.0       5    100
1    6  TSLA     10.0      10    200
2    6  ORCL     10.0      15    300
3    5  MSFT     10.0      20    400
4    5  TSLA     10.0      25    500
5    5  ORCL     10.0      30    600
9    3  MSFT      5.0      50   1000
10   3  TSLA      5.0      55   1100
11   3  ORCL      5.0      60   1200
12   2  MSFT      5.0      65   1300
13   2  TSLA      5.0      70   1400
14   2  ORCL      5.0      75   1500
15   1  MSFT     15.0      80   1600
16   1  TSLA     15.0      85   1700
17   1  ORCL     15.0      90   1800

另一種選擇是使用groupby.head一旦你drop_duplicated保留唯一版本-ID夫婦。 然后在merge使用 select version-id 。

df.merge(df[['version','id']].drop_duplicates().groupby('version').head(n))