Изменение формы фрейма данных groupby до фиксированных размеров

#python #pandas #numpy

#python #pandas #numpy

Вопрос:

У меня есть dataframe df со следующими данными.

 A    B    C    D
1    1    3    1
1    2    9    8
1    3    3    9
2    1    2    9
2    2    1    4
2    3    9    5
2    4    6    4
3    1    4    1
3    2    0    4
4    1    2    6
5    1    2    4
5    2    8    3

grp = df.groupby('A')
  

Затем я хочу, чтобы все группы dataframe df, сгруппированные по столбцам A, имели одинаковое количество строк. Либо обрезать дополнительные строки, либо добавить 0 строк. Для приведенных выше данных я хочу, чтобы все группы имели 3 строки. Мне требовались следующие результаты.

 A    B    C    D
1    1    3    1
1    2    9    8
1    3    3    9
2    1    2    9
2    2    1    4
2    3    9    5
3    1    4    1
3    2    0    4
3    0    0    0
4    1    2    6
4    0    0    0
4    0    0    0
5    1    2    4
5    2    8    3
5    0    0    0
  

Аналогично, я могу захотеть использовать groupby для нескольких столбцов, например
grp = df.groupby([‘A’,’B’])

Ответ №1:

Используйте GroupBy.cumcount для столбца счетчика с DataFrame.reindex by MultiIndex.from_product :

 df['g'] = df.groupby('A').cumcount()
mux = pd.MultiIndex.from_product([df['A'].unique(), range(3)], names=('A','g'))

df = (df.set_index(['A','g'])
        .reindex(mux, fill_value=0)
        .reset_index(level=1, drop=True)
        .reset_index())
print (df)
    A  B  C  D
0   1  1  3  1
1   1  2  9  8
2   1  3  3  9
3   2  1  2  9
4   2  2  1  4
5   2  3  9  5
6   3  1  4  1
7   3  2  0  4
8   3  0  0  0
9   4  1  2  6
10  4  0  0  0
11  4  0  0  0
12  5  1  2  4
13  5  2  8  3
14  5  0  0  0
  

Другое решение с DataFrame.merge левым соединением с помощником DataFrame :

 from  itertools import product

df['g'] = df.groupby('A').cumcount()
df1 = pd.DataFrame(list(product(df['A'].unique(), range(3))), columns=['A','g'])

df = df1.merge(df, how='left').fillna(0).astype(int).drop('g', axis=1)
print (df)
    A  B  C  D
0   1  1  3  1
1   1  2  9  8
2   1  3  3  9
3   2  1  2  9
4   2  2  1  4
5   2  3  9  5
6   3  1  4  1
7   3  2  0  4
8   3  0  0  0
9   4  1  2  6
10  4  0  0  0
11  4  0  0  0
12  5  1  2  4
13  5  2  8  3
14  5  0  0  0
  

Редактировать:

 df['g'] = df.groupby(['A','B']).cumcount()

mux = pd.MultiIndex.from_product([df['A'].unique(), 
                                  df['B'].unique(), 
                                  range(3)], names=('A','B','g'))
df = (df.set_index(['A','B','g'])
        .reindex(mux, fill_value=0)
        .reset_index(level=2, drop=True)
        .reset_index())
print (df.head(10))
   A  B  C  D
0  1  1  3  1
1  1  1  0  0
2  1  1  0  0
3  1  2  9  8
4  1  2  0  0
5  1  2  0  0
6  1  3  3  9
7  1  3  0  0
8  1  3  0  0
9  1  4  0  0
  

 from  itertools import product

df['g'] = df.groupby(['A','B']).cumcount()
df1 = pd.DataFrame(list(product(df['A'].unique(),
                                df['B'].unique(), 
                                range(3))), columns=['A','B','g'])

df = df1.merge(df, how='left').fillna(0).astype(int).drop('g', axis=1)
  

Комментарии:

1. Это отлично работает, если я использую один столбец для groupby. Какие изменения я должен внести в коде, если мне нужно группировать по нескольким столбцам, например. dfg = df.groupby([‘A’,’B’])

2. у них нет группы с ключом (1,4). строки с индексом 9,10,11 не должны отображаться в результате.

3. @AsifKhan — Сначала отфильтруйте идентификатор решения по df = df[df.groupby('A').cumcount() < 3]