如何在 Pandas DataFrame 中的分组变量中选择列的前 N 个键排序值

Question

我有一个数据集：

import pandas as pd

data = [
    ('A', 'X'),
    ('A', 'X'),
    ('A', 'Y'),
    ('A', 'Z'),
    ('B', 1),
    ('B', 1),
    ('B', 2),
    ('B', 2),
    ('B', 3),
    ('B', 3),
    ('C', 'L-7'),
    ('C', 'L-9'),
    ('C', 'L-9'),
    ('T', 2020),
    ('T', 2020),
    ('T', 2025)
]

df = pd.DataFrame(data, columns=['ID', 'SEQ'])
print(df)

我想创建一个关键分组 ID 和 SEQ，以便选择每个 ID 组中每个不同 SEQ 的前 2 行

例如 ID A，按照数据集的顺序有 3 个不同的键“A X”、“A Y”和“A Z”，前两个键是“A X”和“A Y”，因此我必须选择前两行（如果可用）因此

“A X”、“A X”、“A Y”为什么？因为“A Z”是另一个键。

我尝试过使用 groupby 和 head 函数，但我找不到实现此特定结果的方法。有人可以提供解决方案或指出我正确的方向吗？

(df
.groupby(['ID','SEQ'])
.head(2)
)

但是正在重新调整原始数据集，我想知道你们是否可以帮助我使用方法改变来解决这个问题，因为这是我在 pandas 中最喜欢的风格，非常感谢

最终正确的输出是

Answer 1

drop_duplicates

，然后使用

groupby

获取每个“ID”的

head

。然后

merge

与原始 DataFrame 一起保留重复的行。

>>> df.drop_duplicates().groupby("ID").head(2).merge(df)

   ID   SEQ
0   A     X
1   A     X
2   A     Y
3   B     1
4   B     1
5   B     2
6   B     2
7   C   L-7
8   C   L-9
9   C   L-9
10  T  2020
11  T  2020
12  T  2025

Answer 2

使用 groupby 然后使用 head(2) 的方法是正确的，可以获取每个 ID 组内每个不同 SEQ 的前 2 行。

但是，额外的要求是仅获取每个 ID 内的前 2 个唯一 SEQ 组。要实现这一目标，您可以：

创建一个新列，其中具有每个 ID 组内唯一 SEQ 的排名。 使用此排名来过滤数据。 最后，使用原始方法获取每个 ID 组中每个 SEQ 的前 2 行。 这是使用方法链的解决方案：

result = (df
          .assign(rank=df.groupby('ID')['SEQ'].transform(lambda x: x.rank(method='dense')))
          .query('rank <= 2')
          .groupby(['ID', 'SEQ'])
          .head(2)
          .drop(columns=['rank'])
         )

print(result)

这应该会给你想要的输出。

Answer 3

IIUC，您必须仅按 ID 进行分组，然后仅选择唯一的行（例如使用

drop_duplicates

），然后您可以使用

merge

检索这些行：

df = df.\
    merge(
        df.\
        drop_duplicates().\
        groupby(["ID"]).\
        head(2),
        on=["ID", "SEQ"],
        how="right"
    )

df
Out[16]: 
   ID   SEQ
0   A     X
1   A     X
2   A     Y
3   B     1
4   B     1
5   B     2
6   B     2
7   C   L-7
8   C   L-9
9   C   L-9
10  T  2020
11  T  2020
12  T  2025