Pandas: Drop Rows, если размер группы больше среднего

Question

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

d = {'ID':['a12', 'a12','a12','a12','a12','b33','b33','b33','b33','v55','v55','v55','v55','v55','v55'], 'Exp_A':[2.2,2.2,2.2,2.2,2.2,3.1,3.1,3.1,3.1,1.5,1.5,1.5,1.5,1.5,1.5], 
     'Exp_B':[2.4,2.4,2.4,2.4,2.4,1.2,1.2,1.2,1.2,1.5,1.5,1.5,1.5,1.5,1.5], 
     'A':[0,0,1,0,1,0,1,0,1,0,1,1,1,0,1], 'B':[0,0,1,1,1,0,0,1,1,1,0,0,1,0,1]}
df1 = pd.DataFrame(data=d)

Я хочу, чтобы все df1.ID были размером df1.groupby('ID').size().mean().Так что df1 должно выглядеть так:

    A   B   Exp_A   Exp_B   ID
0   0   0   2.2      2.4    a12
1   0   0   2.2      2.4    a12
2   1   1   2.2      2.4    a12
3   0   1   2.2      2.4    a12
4   1   1   2.2      2.4    a12

5   0   0   3.1      1.2    b33
6   1   0   3.1      1.2    b33
7   0   1   3.1      1.2    b33
8   1   1   3.1      1.2    b33
9   0   0   3.1      1.2    b33

10  0   1   1.5      1.5    v55
11  1   0   1.5      1.5    v55
12  1   0   1.5      1.5    v55
13  1   1   1.5      1.5    v55
14  0   0   1.5      1.5    v55

Answer 1

Вот решение без зацикливания.Вы можете сначала определить количество строк для каждого идентификатора, а затем приступить к изменению материала.

# Getting the minimum required number of rows for each ID
min_req = df.groupby('ID').size().mean()

# Adding auto-increment column with respect to ID column
df['row_count'] = df.groupby(['ID']).cumcount()+1

# Adding excess rows equal to required rows
# we will delete unneeded ones later
df2 = df.groupby('ID', as_index=False).max()
df2 = df2.loc[df2['row_count']<int(min_req)]
df2 = df2.assign(A=0, B=0)
df = df.append([df2]*int(min_req), ignore_index=True)

# recalculating the count
df = df.drop('row_count', axis=1)
df = df.sort_values(by=['ID', 'A', 'B'], ascending=[True, False, False])
df['row_count'] = df.groupby(['ID']).cumcount()+1

# Dropping excess rows
df = df.drop((df.loc[df['row_count']>5]).index)
df = df.drop('row_count', axis=1)

df

    A  B  Exp_A  Exp_B   ID
0   0  0    2.2    2.4  a12
1   0  0    2.2    2.4  a12
2   1  1    2.2    2.4  a12
3   0  1    2.2    2.4  a12
4   1  1    2.2    2.4  a12
17  0  0    3.1    1.2  b33
16  0  0    3.1    1.2  b33
15  0  0    3.1    1.2  b33
18  0  0    3.1    1.2  b33
19  0  0    3.1    1.2  b33
10  1  0    1.5    1.5  v55
11  1  0    1.5    1.5  v55
12  1  1    1.5    1.5  v55
13  0  0    1.5    1.5  v55
14  1  1    1.5    1.5  v55

Answer 2

Вот одно решение, использующее GroupBy.Сложность возникает с вашим условием добавления дополнительных строк с определенными столбцами, установленными в 0, всякий раз, когда конкретная группа слишком мала.

g = df1.groupby('ID')
n = int(g.size().mean())

res = []
for _, df in g:
    k = len(df.index)
    excess = n - k
    if excess > 0:
        df = df.append(pd.concat([df.iloc[[-1]].assign(A=0, B=0)]*excess))
    res.append(df.iloc[:n])

res = pd.concat(res, ignore_index=True)

print(res)

    A  B  Exp_A  Exp_B   ID
0   0  0    2.2    2.4  a12
1   0  0    2.2    2.4  a12
2   1  1    2.2    2.4  a12
3   0  1    2.2    2.4  a12
4   1  1    2.2    2.4  a12
5   0  0    3.1    1.2  b33
6   1  0    3.1    1.2  b33
7   0  1    3.1    1.2  b33
8   1  1    3.1    1.2  b33
9   0  0    3.1    1.2  b33
10  0  1    1.5    1.5  v55
11  1  0    1.5    1.5  v55
12  1  0    1.5    1.5  v55
13  1  1    1.5    1.5  v55
14  0  0    1.5    1.5  v55