如何计算熊猫中分组数量的相关系数

Question

我有一个DataFrame，其中每一行代表交通事故。其中两列是 Speed_limit 和 Number_of_casualties 。我想计算限速与每个限速的伤亡人数与事故数之比之间的皮尔逊相关系数。

到目前为止，我的解决方案是将相关数量作为数组获取并使用 SciPy 的pearsonr：

import pandas as pd
import scipy.stats
df = pd.DataFrame({'Speed_limit': [10, 10, 20, 20, 20, 30],
                   'Number_of_casualties': [1, 2, 3, 4, 1, 4]})
accidents_per_speed_limit = df['Speed_limit'].value_counts().sort_index()
number_of_casualties_per_speed_limit = df.groupby('Speed_limit').sum()['Number_of_casualties']
speed_limit = accidents_per_speed_limit.index
ratio = number_of_casualties_per_speed_limit.values / accidents_per_speed_limit.values
r, _ = scipy.stats.pearsonr(x=speed_limit, y=ratio)
print("The Pearson's correlation coefficient between the number of casualties per accidents and the speed limit is {r}.".format(r=r))

然而，在我看来，应该有可能使用熊猫更优雅地做到这一点。DataFrame.corr 方法。如何重构此代码以使其更pandas？

Answer 1

而不是计数和总和，您可以直接使用分组数据mean然后使用series corr（默认情况下方法是pearson），即

m = df.groupby('Speed_limit').mean().reset_index()
m['Speed_limit'].corr(m['Number_of_casualties'])

输出：

0.99926008128973687

Answer 2

我使用两个辅助DataFrames找到了以下方法：

df_aux = df.groupby('Speed_limit').agg(['count', 'sum'])
df_aux2 = pd.DataFrame({'ratio': df_aux['Number_of_casualties', 'sum'] / df_aux['Number_of_casualties', 'count'],
                        'speed_limit': df_aux.index})
print(df_aux2.corr()['ratio']['speed_limit'])

这证实了用scipy.stats.pearsonr获得的结果。不过，它仍然不是很优雅，我将不胜感激改进建议。