python - 获取 pandas DataFrame 中每个组的前 N ​​个最大行

Question

我有一个数据框

val1    val1     distance
100      200      1.5
100      300       5.2
100      234       3.7
200      100       1.6
200      600       4.8

我想找到每个 val1 的前 2 个距离行。 我是

对于每个 val1 获取前 2 个最小距离行

   val1  val2  distance
0   100   200       1.5
2   100   234       3.7
3   200   100       1.6
4   200   600       4.8

这只是一个示例数据集(我实际上有大约 1000 个唯一的 val1) 因此我认为按 val1 分组不是一个好主意。

有没有更快的方法？

Answer 1

看来我过度分析了你的问题，但是一个简单的 sort_values ，后跟 groupby + head 应该可以满足你的需求。

df.sort_values(['val1', 'distance']).groupby('val1').head(2)

   val1  val2  distance
0   100   200       1.5
2   100   234       3.7
3   200   100       1.6
4   200   600       4.8

如果您需要超过 2 个，只需将参数更改为 head 即可。此外，例如，如果您调用 head(2) 并且只有一行，则返回该行(没有任何问题)。