python - Pandas :如果组的大小大于平均值，则删除行

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

这是一个使用 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