首页 > 解决方案 > pandas-groupby:应用需要 2 列作为输入的自定义函数来获得一列作为输出

问题描述

我有一个包含日期和每天值的数据框。我想查看值的梯度,如果它在增长,下降,......最好的方法是应用线性回归,将天作为 x,将值作为 y:

import pandas as pd
df = pd.DataFrame({'customer':['a','a','a','b','b','b'],
                   'day':[1,2,4,2,3,4],
                   'value':[1.5,2.4,3.6,1.5,1.3,1.1]})

东风:


  customer  day  value
0        a    1    1.5
1        a    2    2.4
2        a    4    3.6
3        b    2    1.5
4        b    3    1.3
5        b    4    1.1

手动我可以做一个线性回归:

from sklearn.linear_model import LinearRegression
def gradient(x,y):
    return LinearRegression().fit(x,y).coef_[0]

xa = df[df.customer =='a'].day.values.reshape(-1, 1)
ya = df[df.customer =='a'].value.values.reshape(-1, 1)
xb = df[df.customer =='b'].day.values.reshape(-1, 1)
yb = df[df.customer =='b'].value.values.reshape(-1, 1)
print(gradient(xa,ya),gradient(xb,yb))

result: [0.68571429] [-0.2]

但我想使用groupby

df.groupby('customer').agg({'value':['mean','sum','gradient']})

输出如下:

               value     
          mean  sum  gradient
customer           
a          2.5  7.5   0.685
b          1.3  3.9  -0.2

问题是梯度需要 2 列作为输入。

标签: pandaspandas-groupby

解决方案


你可以做:

# calculate gradient
v =  (df
      .groupby('customer')
      .apply(lambda x: gradient(x['day'].to_numpy().reshape(-1, 1), 
                                x['value'].to_numpy().reshape(-1, 1)))
v.name = 'gradient'

# calculate mean, sum
d1 = df.groupby('customer').agg({'value': ['mean', 'sum']})

# join the results
d1 = d1.join(v)

# fix columns
d1.columns = d1.columns.str.join('')
print(d1)

               valuemean     valuesum  gradient
customer                                       
a                   2.5           7.5  0.685714
b                   1.3           3.9 -0.200000

推荐阅读