如何通过聚合函数筛选总库存>0的Pandas DataFrame店铺记录？

需求说明 #

从给定的DataFrame中，保留所有总库存大于0的店铺的全部记录，排除总库存为0的店铺（如示例中的store2）。

实现方法 #

方法一：使用groupby + transform直接标记筛选

这种方法无需额外存储店铺列表，直接通过transform生成与原数据同长度的判断序列，一步完成筛选：

import pandas as pd

# 构造示例DataFrame
data = {
    'STORE': ['store1', 'store1', 'store2', 'store2', 'store3', 'store3'],
    'PRODUCT': ['a', 'b', 'a', 'b', 'a', 'b'],
    'INVENTORY': [1, 0, 0, 0, 1, 1]
}
df = pd.DataFrame(data, index=[1,2,3,4,5,6])

# 生成每个店铺的总库存判断序列
has_positive_inventory = df.groupby('STORE')['INVENTORY'].transform('sum') > 0

# 筛选符合条件的记录
filtered_df = df[has_positive_inventory]
print(filtered_df)

方法二：先获取有效店铺列表再筛选

如果需要单独保存有效店铺的列表，可以先计算总库存筛选出店铺，再用isin匹配：

import pandas as pd

# 构造示例DataFrame
data = {
    'STORE': ['store1', 'store1', 'store2', 'store2', 'store3', 'store3'],
    'PRODUCT': ['a', 'b', 'a', 'b', 'a', 'b'],
    'INVENTORY': [1, 0, 0, 0, 1, 1]
}
df = pd.DataFrame(data, index=[1,2,3,4,5,6])

# 筛选出总库存>0的店铺
valid_stores = df.groupby('STORE')['INVENTORY'].sum().loc[lambda x: x > 0].index

# 匹配筛选记录
filtered_df = df[df['STORE'].isin(valid_stores)]
print(filtered_df)

两种方法最终都会输出目标结果：

STORE PRODUCT  INVENTORY
1  store1       a          1
2  store1       b          0
5  store3       a          1
6  store3       b          1

内容的提问来源于stack exchange，提问作者Ari Frid