当N大于组数时nlargest(N)的行为?

时间:2019-02-08 11:23:02

标签: python pandas pandas-groupby

我从以下列表中构建了一个DataFrame

df_list_1 = [{"animal": "dog", "color": "red", "age": 4, "n_legs": 4,}, 
             {"animal": "dog", "color": "blue", "age": 4, "n_legs": 3},
             {"animal": "cat", "color": "blue", "age": 4, "n_legs": 4},
             {"animal": "dog", "color": "yellow", "age": 5, "n_legs":2},
             {"animal": "dog", "color": "white", "age": 4, "n_legs": 2},
             {"animal": "dog", "color": "black", "age": 4, "n_legs": 4},
             {"animal": "cat", "color": "brown", "age": 4, "n_legs": 4}]

我现在想获得一个新的数据框,该数据框仅显示每个组中具有相同的age的前4个条目(按n_legs排序)。

为此,我尝试了

dfg = df_1.set_index(["animal", 'color']).groupby("n_legs")['age'].nlargest(4).reset_index()

但是,这给了我一个数据帧,其中n_legs列已删除。

    animal  color   age
0   dog     red     4
1   dog     blue    4
2   cat     blue    4
3   dog     yellow  5
4   dog     white   4
5   dog     black   4
6   cat     brown   4

我猜这是因为4等于最大组中的元素数。实际上,如果我这样做

dfg = df_1.set_index(["animal", 'color']).groupby("n_legs")['age'].nlargest(3).reset_index()

我得到以下

    n_legs  animal  color   age
0   2       dog     yellow  5
1   2       dog     white   4
2   3       dog     blue    4
3   4       dog     red     4
4   4       cat     blue    4
5   4       dog     black   4

这是预期的行为吗?

即使使用nlargest(N)N大于最大组中元素的数量,有没有办法始终显示列?

谢谢!

1 个答案:

答案 0 :(得分:2)

我认为是bug 16345

替代解决方案运行良好,速度明显更快-首先sort_values,然后调用GroupBy.head

dfg = (df_1.sort_values(["animal", 'color','age'], ascending=[False, False, True])
          .groupby("n_legs")
          .head(4))