Create new column on grouped data frame

Question

I want to create new column that is calculated by groups using multiple columns from current data frame. Basically something like this in 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)

In pandas I think I should use groupby and apply to create new column and then assign it to the original data frame. This is what I've tried so far:

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

What is preferred way (ideally without sorting the data) to achieve this?

Answer 1

Add parameter group_keys=False for avoid MultiIndex , so possible assign back to new column:

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

Alternative is remove first level:

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

Detail :

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

Now you can also implement it in python with datar in the way exactly you did in 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)]

I am the author of the package. Feel free to submit issues if you have any questions.