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

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