在分组的数据框中创建新列

时间:2019-01-10 09:16:36

标签: python pandas

我想创建一个新列,该列是使用来自当前数据帧的多个列按组计算的。在Rtidyverse)中基本上是这样的:

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中,我应该使用groupbyapply创建新列,然后将其分配给原始数据帧。到目前为止,这是我尝试过的:

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

实现此目标的首选方式是什么(理想情况下无需对数据进行排序)?

2 个答案:

答案 0 :(得分: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

答案 1 :(得分:1)

现在你也可以用 datar 在 python 中实现它,就像你在 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)]

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