Pandas 按唯一计数分组作为新列

我想添加一个新列col在我的 pandas 数据框中，计算公式为：

select count(distinct ITEM) as col
from base_data
where STOCK > 0 
group by DEPT, CLAS, DATE;

我正在做的事

assort_size = base_data[(base_data['STOCK'] > 0)]\
.groupby(['DEPT','CLAS','DATE'])['ITEM']\
.transform('nunique')

基本上对于每个部门、班级、日期组合，我想获取库存中的商品数量。所以我想要与父数据框合并的结果，但结果如下pandas.core.series.Series因此我不能append (axis=1)它返回（行数不同，例如 1.6 M 与 1.4 M）。还有我没有DEPT, CLAS, DATE要加入的列。我可以在这里做什么来获取按列分组的数据框？

有没有更好的方法直接在父pandas数据框中创建新列（base_data）而不是像我一样创建一个新对象assort_size?

您可以使用boolean indexing首先，然后groupby with nunique最后join:

base_data = pd.DataFrame({"DEPT": ["a", "a", "b", "b"],
                   "CLAS":['d','d','d','d'],
                   "STOCK": [-1, 1, 2,2],
                   "DATE":pd.to_datetime(['2001-10-10','2001-10-10',
                                          '2001-10-10','2001-10-10']),
                   "ITEM":[1,2,3,4]})

print (base_data)
  CLAS       DATE DEPT  ITEM  STOCK
0    d 2001-10-10    a     1     -1
1    d 2001-10-10    a     2      1
2    d 2001-10-10    b     3      2
3    d 2001-10-10    b     4      2

assort_size = base_data[(base_data['STOCK'] > 0)]\
.groupby(['DEPT','CLAS','DATE'])['ITEM'].nunique().rename('n_item')
print (assort_size)
DEPT  CLAS  DATE      
a     d     2001-10-10    1
b     d     2001-10-10    2
Name: n_item, dtype: int64

print (base_data.join(assort_size, on=['DEPT','CLAS','DATE']))
  CLAS       DATE DEPT  ITEM  STOCK  n_item
0    d 2001-10-10    a     1     -1       1
1    d 2001-10-10    a     2      1       1
2    d 2001-10-10    b     3      2       2
3    d 2001-10-10    b     4      2       2