Создать новый столбец в сгруппированном фрейме данных

1

Я хочу создать новый столбец, который рассчитывается по группам, используя несколько столбцов из текущего фрейма данных. В основном как-то так в 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

Какой предпочтительный способ (в идеале без сортировки данных) для достижения этой цели?

Теги:
pandas

1 ответ

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
    Это именно то, что я искал. Спасибо

Ещё вопросы

Сообщество Overcoder
Наверх
Меню