如何筛选Pandas DataFrame中对应单一公司的州的行及获取此类州集合？

1. 生成目标布尔数组 #

先统计每个州对应的唯一公司数量，再将该信息映射到原DataFrame的每一行，判断该行所在州的唯一公司数是否为1：

import pandas as pd

df = pd.DataFrame({"product": [1,2,3,4,5], "company": ["A", "B","B","A","B"], "state": ["CA", "NY", "CA", "CA", "NY"]})

# 计算每个州的唯一公司数量
state_unique_companies = df.groupby('state')['company'].nunique()

# 生成布尔数组，标记所在州仅对应单一公司的行
bool_array = df['state'].map(state_unique_companies) == 1
print(bool_array.values)  # 输出：[False  True False False  True]

2. 获取仅含单一公司的州的集合 #

方法一：直接从统计结果筛选 #

从之前的统计结果中提取唯一公司数为1的州，转成集合：

single_company_states = set(state_unique_companies[state_unique_companies == 1].index)
print(single_company_states)  # 输出：{'NY'}

方法二：利用布尔数组筛选 #

如果已经生成布尔数组，可通过布尔索引筛选目标行，再提取州列去重后转集合：

single_company_states = set(df[bool_array]['state'].unique())
print(single_company_states)  # 输出：{'NY'}

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