Как сгруппировать по столбцу в pandas и применить ifelse на основе значений столбца

Question

Я пытаюсь сгруппировать по столбцу в панде DataFrame и применить функцию ifelse, у которой есть тест, который использует значения столбцов и возвращает разные значения для true и false в зависимости от теста.

Я легко выполнил это в R с помощью пакета data.table, и я бы хотел, чтобы в Python было то же самое, используя панд.

library(data.table)

dat <- data.table(
  i_n = c('a', 'b', 'b', 'd', 'e', 'f', 'g', 'h', 'h', 'i'),
  m_b_r = c(0, 1, 0, 0, 0, 0, 0, 1, 0, 0),
  o_q = c(1, 8, 8, 1, 1, 1, 1, 2, 2, 1),
  a_q = c(1, 5, 15, 1, 57, 1, 5, 1, 1, 1)
)

dat <- dat[, ('p_q') := ifelse(
  test =get('m_b_r') == 1,
  yes =get('a_q'),
  no =get('o_q') - data.table::shift(get('a_q'), n = 1L, type = 'lag', fill = 0 )
), by = 'i_n']

Answer 1

Я считаю, что нужно numpy.where с DataFrameGroupBy.shift:

shifted = dat.groupby('i_n')['a_q'].shift().fillna(0)
dat['p_q'] = np.where(dat['m_b_r'] == 1, dat['a_q'], dat['o_q'] - shifted)
print (dat)
  i_n  m_b_r  o_q  a_q  p_q
0   a      0    1    1  1.0
1   b      1    8    5  5.0
2   b      0    8   15  3.0
3   d      0    1    1  1.0
4   e      0    1   57  1.0
5   f      0    1    1  1.0
6   g      0    1    5  1.0
7   h      1    2    1  1.0
8   h      0    2    1  1.0
9   i      0    1    1  1.0

Возможен вызов пользовательской функции, но первое решение быстрее:

def f(x):
    x['p_q'] = np.where(x['m_b_r'] == 1, x['a_q'], x['o_q'] - x['a_q'].shift().fillna(0))
    return x

df = dat.groupby('i_n').apply(f)
print (df)
  i_n  m_b_r  o_q  a_q  p_q
0   a      0    1    1  1.0
1   b      1    8    5  5.0
2   b      0    8   15  3.0
3   d      0    1    1  1.0
4   e      0    1   57  1.0
5   f      0    1    1  1.0
6   g      0    1    5  1.0
7   h      1    2    1  1.0
8   h      0    2    1  1.0
9   i      0    1    1  1.0