如何以特殊方式对熊猫数据框进行排序

Question

给出一个熊猫数据框

df = pd.DataFrame({'a': [1,2,3,4,5,6,7,8],
                   'b': [0,0,1,1,2,2,3,3]})

如何沿b列对它进行排序，以使其重新排列为{0,1,2,3,0,1,2,3}。

即结果数据框是

1   0
3   1
5   2
7   3
2   0
4   1
6   2
8   3

Answer 1

使用cumcount

添加一列

df.assign(x=df.groupby('b').cumcount()).sort_values(['x', 'b']).drop('x', axis=1)

   a  b
0  1  0
2  3  1
4  5  2
6  7  3
1  2  0
3  4  1
5  6  2
7  8  3

---

Numpy的lexsort，iloc和cumcount

df.iloc[np.lexsort([df['b'], df.groupby('b').cumcount()])]

   a  b
0  1  0
2  3  1
4  5  2
6  7  3
1  2  0
3  4  1
5  6  2
7  8  3

Answer 2

让我们尝试一下：

s = df.groupby('b').cumcount().sort_values(kind='mergesort')

df = df.loc[s.index]

输出：

   a  b
0  1  0
2  3  1
4  5  2
6  7  3
1  2  0
3  4  1
5  6  2
7  8  3

Answer 3

具有快速基准测试的另一种解决方案：

pd.DataFrame(sorted(df.values.tolist(), key=lambda k: (k[0]%2==0, k[1])), columns=['a', 'b'])

打印：

   a  b
0  1  0
1  3  1
2  5  2
3  7  3
4  2  0
5  4  1
6  6  2
7  8  3

基准：

df = pd.DataFrame({'a': [1,2,3,4,5,6,7,8],
                   'b': [0,0,1,1,2,2,3,3]})


from timeit import timeit

def f1():
    return pd.DataFrame(sorted(df.values.tolist(), key=lambda k: (k[0]%2==0, k[1])), columns=['a', 'b'])

def f2():
    s = df.groupby('b').cumcount().sort_values(kind='mergesort')
    return df.loc[s.index]

def f3():
    return df.iloc[np.lexsort([df['b'], df.groupby('b').cumcount()])]

t1 = timeit(lambda: f1(), number=1_000)
t2 = timeit(lambda: f2(), number=1_000)
t3 = timeit(lambda: f3(), number=1_000)

print(t1)
print(t2)
print(t3)

图片（AMD 2400G / Ubuntu 18.04，Python 3.8.3，Pandas 1.0.3）：

0.45131446300365496
2.2533202580088982
1.6977271080104401

因此，使用sorted()的解决方案似乎最快。

编辑：在大约3000个元素之后，np.lexsort的解决方案开始获胜。

Answer 4

这是针对您问题的简单化处理，它基于b列已被排序的知识：

res = df.to_numpy()

pd.DataFrame(np.vstack((res[::2],res[1::2])),columns=df.columns)

#similar
pd.concat((df.iloc[::2], df.iloc[1::2]))


    a   b
0   1   0
1   3   1
2   5   2
3   7   3
4   2   0
5   4   1
6   6   2
7   8   3