новый столбец с фреймом

Question

import pandas
import numpy

df = pandas.DataFrame({'id_1' : [1,2,1,1,1,1,1,2,2,2,2], 
                      'id_2' :  [1,1,1,1,1,2,2,2,2,2,2],
                      'v_1' :   [2,1,1,3,2,1,2,4,1,1,2],
                      'v_2' :   [1,1,1,1,2,2,2,1,1,2,2],
                      'v_3' :   [3,3,3,3,4,4,4,3,3,3,3]})


In [4]: df                                                                                                                                                                                                  
Out[4]: 
    id_1  id_2  v_1  v_2  v_3
0      1     1    2    1    3
1      2     1    1    1    3
2      1     1    1    1    3
3      1     1    3    1    3
4      1     1    2    2    4
5      1     2    1    2    4
6      1     2    2    2    4
7      2     2    4    1    3
8      2     2    1    1    3
9      2     2    1    2    3
10     2     2    2    2    3

sub = df[(df['id_1'] == 1) & (df['id_2'] == 1)].copy()
sub['v_4'] = numpy.where(sub['v_1'] == sub['v_2'].shift(), 'A', \
                         numpy.where(sub['v_1'] == sub['v_3'].shift(), 'B', 'C'))


In [6]: sub                                                                                                                                                                                                 
Out[6]: 
   id_1  id_2  v_1  v_2  v_3 v_4
0     1     1    2    1    3   C
2     1     1    1    1    3   A
3     1     1    3    1    3   B
4     1     1    2    2    4   C

У меня есть кадр данных, как определено выше. Я хотел бы выполнить некоторую операцию, в основном классифицировать, равняется ли v_1 предыдущему v_2 или v_3 для каждой группы (id_1, id_2). Я выполнил операцию, которая выполняется на sub df. И я хотел бы иметь однострочный код для объединения следующей группы вместе с операцией, которую я выполняю на sub df вместе.

gbdf = df.groupby(by=['id_1', 'id_2'])

Я пробовал что-то вроде

gbdf['v_4'] = numpy.where(gbdf['v_1'] == gbdf['v_2'].shift(), 'A', \
                         numpy.where(gbdf['v_1'] == gbdf['v_3'].shift(), 'B', 'C'))

и ошибка была

'DataFrameGroupBy' object does not support item assignment

Я также пытался

df['v_4'] = numpy.where(gbdf['v_1'] == gbdf['v_2'].shift(), 'A', \
                         numpy.where(gbdf['v_1'] == gbdf['v_3'].shift(), 'B', 'C'))

, который, как я считаю, был неверным результатом, он не выравнивает результат группировки по исходному порядку.

Мне интересно, есть ли изящный способ достичь этого.

Answer 1

Появляется список фреймов данных, соответствующих содержимому фрейма данных sub, но для всех результатов .groupby():

import numpy
import pandas

source = pandas.DataFrame(
    {'id_1': [1, 2, 1, 1, 1, 1, 1, 2, 2, 2, 2],
     'id_2': [1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2],
     'v_1': [2, 1, 1, 3, 2, 1, 2, 4, 1, 1, 2],
     'v_2': [1, 1, 1, 1, 2, 2, 2, 1, 1, 2, 2],
     'v_3': [3, 3, 3, 3, 4, 4, 4, 3, 3, 3, 3]})


def add_v4(df):
    df['v_4'] = numpy.where(df['v_1'] == df['v_2'].shift(), 'A', numpy.where(df['v_1'] == df['v_3'].shift(), 'B', 'C'))
    return df


dfs = [add_v4(pandas.DataFrame(slice)) for _, slice in source.groupby(by=['id_1', 'id_2'])]
print(dfs)

Об этой строке:

dfs = [add_v4(pandas.DataFrame(slice)) for _, slice in source.groupby(by=['id_1', 'id_2'])]

Это понимание списка, которое получает все фрагменты из groupby и превращает их в фактические новые кадры данных перед передачей их в add_v4, что возвращает измененный кадр данных, который будет добавлен в список.