如何将 DataFrame 行数限制为特定列中的第 X 个唯一值？

Question

例如，我们有以下 DataFrame：

A B
1 2
1 2
2 3
3 4
4 5 
4 2

我们会知道我们想要在 A 列中有 x（比如 3）个唯一值。那么所需的 output 将是：

A B
1 2
1 2
2 3
3 4

我考虑过遍历有问题的列，通过跟踪并获取具有正确索引的 DataFrame 的子集来计算唯一值的数量。 我仍然是 Python 的新手，我相信会有更有效的方法来做到这一点，请分享您的解决方案。 赞赏！

Answer 1

您可以尝试series.factorize索引从 0 开始的唯一值，然后 select 是 <= n-1 的值（因为索引从 0 开始），因此也保留订单：

n=3
df[df['A'].factorize()[0]<=n-1]

---

   A  B
0  1  2
1  1  2
2  2  3
3  3  4

Answer 2

您可以使用np.random.choice到 select 唯一的 id，然后使用这些 id 到isin行：

selected_ids = np.random.choice(df['A'].unique(), replace=False, size=3)

df[df['A'].isin(selected_ids)]