Как особым образом отсортировать pandas фрейм данных

Question

Учитывая pandas фрейм данных

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

Попробуем:

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 2

Добавьте столбец, используя 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 s 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 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

Это упрощенно c и указывает c на ваш вопрос и основывается на знании того, что столбец 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