pandas groupby по нескольким колонкам + фильтрация

Игрушечный датасет для примера.

import pandas as pd
import numpy as np

n = 30 
data = pd.DataFrame({ 
        'column1': np.random.randint(1, 10, size=n),     
        'column2': np.random.randint(1, 10, size=n),  
        'column3': np.random.randint(1, 4, size=n)
        })


    column1 column2 column3
0   2       7       1
1   5       1       1
2   9       9       3
3   2       3       2
4   4       6       3

Нужно получить новый столбец column4, сгруппированный по column1 и column2, и с условием column3 == 1.

Если делаю с группировкой по одной колонке - то всё ок: data.column1.map( train.groupby(train['column1'][train['column3'] == 1]).size()).fillna(0)

Пытался применять filter, loc, не получается.


Ответы (2 шт):

Автор решения: strawdog

Комментариев к вопросу слишком много. попробуем пошагово решать проблему в ответе:

При:

import pandas as pd
import numpy as np

np.random.seed(42)
n = 30
data = pd.DataFrame({'column1': np.random.randint(1, 10, size=n),
                     'column2': np.random.randint(1, 10, size=n),
                     'column3': np.random.randint(0, 4, size=n),
                     'column4': np.random.randint(1, 10, size=n),
                     'column5': np.random.randint(1, 10, size=n),
                     'column6': np.random.randint(1, 10, size=n),
                     'column7': np.random.randint(1, 10, size=n)})

у меня отлично работают выборки:

data.loc[data['column3']!=0, ['column1', 'column2', 'column3']]

и

data.loc[data['column3']==0, ['column1', 'column2', 'column3']]

теперь можно делать нужную вам группировку:

data['res'] = data.loc[data['column3']==0, ['column1', 'column2', 'column3']].groupby(['column1', 'column2']).transform('count')
    column1  column2  column3  column4  column5  column6  column7  res
0         7        3        1        8        7        7        1  NaN
1         4        7        2        8        7        6        3  NaN
2         8        5        3        3        9        8        2  NaN
3         5        9        0        1        3        9        5  1.0
4         7        7        1        8        7        5        6  NaN
5         3        2        3        3        1        1        7  NaN
6         7        4        0        3        4        3        4  1.0
7         8        9        3        1        4        8        7  NaN
и т.д.

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

data['res'] = data['res'].fillna(method="ffill").fillna(method="bfill")
    column1  column2  column3  column4  column5  column6  column7  res
0         7        3        1        8        7        7        1  1.0
1         4        7        2        8        7        6        3  1.0
2         8        5        3        3        9        8        2  1.0
3         5        9        0        1        3        9        5  1.0
4         7        7        1        8        7        5        6  1.0
5         3        2        3        3        1        1        7  1.0
6         7        4        0        3        4        3        4  1.0
7         8        9        3        1        4        8        7  1.0
8         5        2        0        5        5        6        8  1.0
9         4        9        1        7        7        8        1  1.0
10        8        5        2        9        7        9        6  1.0
11        8        2        0        7        4        4        8  1.0
12        3        4        3        9        7        1        5  1.0
13        6        7        1        8        3        1        4  1.0
14        5        8        0        2        6        4        2  2.0
15        2        3        3        1        2        7        6  2.0
16        8        1        3        7        9        2        6  2.0
17        6        4        3        7        5        3        1  2.0
18        2        2        0        8        6        1        9  1.0
19        5        8        0        5        4        5        6  2.0
20        1        4        0        3        7        1        3  1.0
21        6        2        2        8        9        8        4  1.0
и т. д.

то есть, сначала делаем forward fill, затем backward fill. можно и наоборот, но тогда получится другая колонка.

→ Ссылка
Автор решения: Garp

Итого

data['res2'] = ''

data['res1'] = data[['col1', 'col2', 'col3']].loc[data['col3'] == 0].groupby(['col1', 'col2']).transform('count') 

for i in range(0, data.shape[0], 1):
        if data['res1'][i] in range(1, 100, 1): 
            for j in range(0, data.shape[0], 1):
                if data['col1'][j] == data['col1'][i] and data['col2'][j] == data['col2'][i]:
                    data['res2'][j] = data['res1'][i]  

data['res2'] = data['res2'].replace('', 0)   

Очень плохой код, знаю(( Но всё работает

→ Ссылка