如何单次扫描Pandas DataFrame生成两类列表及按分组提取列表？

希望基于另一列的条件，从某一列生成两个列表。目前通过两次扫描Pandas DataFrame实现，但有两个疑问：

1. 是否可以通过单次扫描DataFrame生成这两个列表？
2. 如何按分组提取对应的列表？

原实现代码：

data = {
    "co2": [95, 90, 99, 104, 105, 94, 99, 104],
    "model": [
        "Citigo",
        "Fabia",
        "Fiesta",
        "Rapid",
        "Focus",
        "Mondeo",
        "Octavia",
        "B-Max",
    ],
    "car": ["Skoda", "Skoda", "Ford", "Skoda", "Ford", "Ford", "BMW", "Ford"],
}

df = pd.DataFrame(data)

# 生成两个列表
list_skoda = df.loc[df["car"] == "Skoda", "model"].tolist()
print(f"{list_skoda=}")

list_others = df.loc[df["car"] != "Skoda", "model"].tolist()
print(f"{list_others=}")

# 尝试分组提取
df.groupby(["car"]).apply(print)
l = df.groupby(["car"])["model"].groups

print(f"{l=}") # 仅返回索引而非模型名称

1. 单次扫描生成两个列表 #

可以通过一次遍历或向量化操作实现，避免重复扫描DataFrame，以下是两种高效方式：

方式一：逐行遍历（适合小数据集） #

一次循环完成分类，直接将模型名加入对应列表：

list_skoda = []
list_others = []

for _, row in df.iterrows():
    if row["car"] == "Skoda":
        list_skoda.append(row["model"])
    else:
        list_others.append(row["model"])

print(f"{list_skoda=}")
print(f"{list_others=}")

方式二：向量化分组（适合大数据集） #

利用groupby一次完成分组，再拆分出目标列表，本质是单次扫描：

# 按car分组并将每组model转为列表
groups = df.groupby("car")["model"].apply(list)
list_skoda = groups.get("Skoda", [])
# 合并非Skoda的所有组
list_others = [model for car, models in groups.items() if car != "Skoda" for model in models]

print(f"{list_skoda=}")
print(f"{list_others=}")

或者用掩码一次性计算后拆分，Pandas内部会基于同一掩码操作，实际仅扫描一次：

mask = df["car"] == "Skoda"
list_skoda = df.loc[mask, "model"].tolist()
list_others = df.loc[~mask, "model"].tolist()

2. 按分组提取对应列表 #

原代码返回索引是因为用了.groups，要获取模型名称列表，可直接用apply(list)或agg(list)：

方式一：生成分组字典 #

将分组结果转为字典，键为品牌名，值为对应模型列表：

car_model_dict = df.groupby("car")["model"].apply(list).to_dict()
print(car_model_dict)
# 输出：{'BMW': ['Octavia'], 'Ford': ['Fiesta', 'Focus', 'Mondeo', 'B-Max'], 'Skoda': ['Citigo', 'Fabia', 'Rapid']}

# 提取单个分组
bmw_models = car_model_dict.get("BMW", [])
ford_models = car_model_dict.get("Ford", [])

方式二：遍历分组对象 #

直接遍历分组结果，逐个获取每个品牌的模型列表：

for car_name, group in df.groupby("car"):
    model_list = group["model"].tolist()
    print(f"{car_name}: {model_list}")

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