根据特定列的值计数添加/删除行

Question

鉴于以下数据框

import pandas as pd

df = pd.DataFrame({
    "ID": [ "1", "1", "1", "1", "1", "2", "2", "3", "3", "3", "4", "4"],
    "Feature": [ 2, 6, 4, 5, 6, 3, 1, 6, 3, 5, 7, 1]
})

如果某个 ID 的值计数小于 4（例如 ID 2、3、4 的值计数分别为 2、3、2），则复制该 ID 出现的最后一行并填充它该ID的值计数变为4， 如果某个 ID 的值计数大于 4（例如 ID 1 的值计数为 5），则删除该 ID 出现的最后一行，使该 ID 的值计数变为 4。

所以给定上面的数据框：

          ID Feature
0         1  2
1         1  6
2         1  4
3         1  5
4         1  6
5         2  3
6         2  1
7         3  6
8         3  3
9         3  5
10        4  7
11        4  1

结果应该是

          ID Feature
0         1  2
1         1  6
2         1  4
3         1  5
4         2  3
5         2  1
6         2  1
7         2  1
8         3  6
9         3  3
10        3  5
11        3  5
12        4  7
13        4  1
14        4  1
15        4  1

有没有任何有效的方式/矢量化方式来做到这一点？谢谢

Answer 1

使用DataFrame.reindex 和method='ffill' 进行过滤并将值添加到MultiIndex.from_product 创建的MultiIndex 和GroupBy.cumcount 的计数器，最后删除MultiIndex 中的辅助级别：

df['g'] = df.groupby('ID').cumcount()

mux = pd.MultiIndex.from_product([df['ID'].unique(), range(4)], names=['ID','g'])
df = (df.set_index(['ID','g'])
        .reindex(mux, method='ffill')
        .reset_index(level=1, drop=True)
        .reset_index())
print (df)
   ID  Feature
0   1        2
1   1        6
2   1        4
3   1        5
4   2        3
5   2        1
6   2        1
7   2        1
8   3        6
9   3        3
10  3        5
11  3        5
12  4        7
13  4        1
14  4        1
15  4        1

Answer 2

你可以craft a MultiIndex from a product和reindex，这会将行数限制为4，然后ffill添加缺失值：

idx = pd.MultiIndex.from_product([df['ID'].unique(), range(4)], names=['ID', ''])

(df.set_index(['ID', df.groupby('ID').cumcount()])
   .reindex(idx)
   .ffill(downcast='infer')
   .droplevel(1)
   .reset_index()
)

输出：

   ID  Feature
0   1        2
1   1        6
2   1        4
3   1        5
4   2        3
5   2        1
6   2        1
7   2        1
8   3        6
9   3        3
10  3        5
11  3        5
12  4        7
13  4        1
14  4        1