Могу ли я ускорить групповую работу на многоколонках с последующей разницей с 1 периодом?

Question

У меня есть такой DataFrame.

df = pd.DataFrame({'a': np.random.randint(0, 3, 10**7), 
                   'b': np.random.randint(0, 4000, 10**7), 
                   'c': np.random.random(10**7)}

Я хочу сгруппировать по первым двум столбцам (неупорядоченно, примерно с таким же количеством групп, как я указал здесь), а затем изменить третий столбец.

df.groupby(['a', 'b']).c.diff(periods=-1)

К сожалению, это довольно медленно для меня, но я предполагаю, что это довольно трудоемкая работа.

>>> %timeit df.groupby(['a', 'b']).c.diff(periods=-1)    
10.2 s ± 75.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Есть ли способ избежать группы Панд, чтобы ускорить это? вверх? Я знаю, что когда в Pandas выполняется groupby-> Reduce, его часто можно заменить чистым NumPy, но я не уверен, что есть что-нибудь умное, что можно сделать вместо операции groupby.

Answer 1

Это действительно интересная проблема, так как обычно мы просто выбираем groupby(), хотя это довольно дорого.

Вот еще один подход, который будет быстрее:

1. сортировать по['a','b']
2. вычислить разность, замаскировать последнюю строку в каждом блоке с помощью NaN

и кода:

s = df.sort_values(['a', 'b'])
df['diff_c'] = (s['c'].diff(periods=-1)
                       .where(s.duplicated(['a','b'], keep='last'))
                )

На моем компьютере ваш кодработает в 8,5 с, в то время как выше работает в 1,8 с.