在分组数据框上创建新列答案

【问题标题】：Create new column on grouped data frame在分组数据框上创建新列
【发布时间】：2021-06-09 04:52:38
【问题描述】：

我想使用当前数据框中的多个列创建由组计算的新列。在R (tidyverse) 中基本上是这样的：

require(tidyverse)

data <- data_frame(
  a = c(1, 2, 1, 2, 3, 1, 2),
  b = c(1, 1, 1, 1, 1, 1, 1),
  c = c(1, 0, 1, 1, 0, 0, 1),
)

data %>% 
  group_by(a) %>% 
  mutate(d = cumsum(b) * c)

在pandas 我想我应该使用groupby 和apply 创建新列，然后将其分配给原始数据框。这是我迄今为止尝试过的：

import numpy as np
import pandas as pd

def create_new_column(data):
    return np.cumsum(data['b']) * data['c']    

data = pd.DataFrame({
    'a': [1, 2, 1, 2, 3, 1, 2],
    'b': [1, 1, 1, 1, 1, 1, 1],
    'c': [1, 0, 1, 1, 0, 0, 1],
})

# assign - throws error
data['d'] = data.groupby('a').apply(create_new_column)

# assign without index - incorrect order in output
data['d'] = data.groupby('a').apply(create_new_column).values

# assign to sorted data frame
data_sorted = data.sort_values('a')
data_sorted['d'] = data_sorted.groupby('a').apply(create_new_column).values

实现此目的的首选方法是什么（最好不对数据进行排序）？

【问题讨论】：

标签： python pandas

【解决方案1】：

添加参数group_keys=False 以避免MultiIndex，因此可以分配回新列：

data['d'] = data.groupby('a', group_keys=False).apply(create_new_column)

替代方法是删除第一级：

data['d'] = data.groupby('a').apply(create_new_column).reset_index(level=0, drop=True)

print (data)
   a  b  c  d
0  1  1  1  1
1  2  1  0  0
2  1  1  1  2
3  2  1  1  2
4  3  1  0  0
5  1  1  0  0
6  2  1  1  3

详情：

print (data.groupby('a').apply(create_new_column))
a   
1  0    1
   2    2
   5    0
2  1    0
   3    2
   6    3
3  4    0
dtype: int64

print (data.groupby('a', group_keys=False).apply(create_new_column))
0    1
2    2
5    0
1    0
3    2
6    3
4    0
dtype: int64

【讨论】：

这正是我想要的。谢谢

【解决方案2】：

现在你也可以在 python 中使用datar 实现它，就像你在 R 中所做的那样：

>>> from datar.all import c, f, tibble, cumsum
>>> 
>>> data = tibble(
...   a = c(1, 2, 1, 2, 3, 1, 2),
...   b = c(1, 1, 1, 1, 1, 1, 1),
...   c = c(1, 0, 1, 1, 0, 0, 1),
... )
>>> 
>>> (data >>
...  group_by(f.a) >>
...  mutate(d=cumsum(f.b) * f.c))
   a  b  c  d
0  1  1  1  1
1  2  1  0  0
2  1  1  1  2
3  2  1  1  2
4  3  1  0  0
5  1  1  0  0
6  2  1  1  3
[Groups: ['a'] (n=3)]

我是包的作者。如果您有任何问题，请随时提交问题。

【讨论】：