在分组数据框上创建新列

Question

我想使用当前数据框中的多列创建由组计算的新列。 在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

实现这一目标的首选方法是什么（理想情况下不对数据进行排序）？

Answer 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

Answer 2

现在，您还可以像在 R 中一样，使用datar在 python 中实现它：

>>> 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)]

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