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

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，如果某个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

您可以从产品和reindex中创建MultiIndex，这将限制行数为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