如何用不同元素匹配数据集中的多个元素？ （例如年份和公司）

Question

我正在处理由 CSV 文件组成的数据集，我的任务是对所有特定数据点进行求和，以便将它们视为一个而不是多个。例如，每条记录由（年份、名称、ID、值）组成，其中 ID、名称和年份可以重复。我的最终目标是一个程序，它可以有效地查看每年，然后查看该年中的每个名称，最后对任何相同的 ID 求和。我正在使用 Python 内置的 CSV 阅读器，但也可以使用 Pandas，因为所有 CSV 都将在项目结束时使用它转换为 .dta 文件。我知道 for 循环会适用于此，我只需检查每一行是否有重复的年份、名称和 ID，但是我正在处理超过一百万条记录，因此优化会有很大帮助。

Answer 1

import pandas as pd

# import csv file
#df = pd.read_csv('test.csv')

# Toy dataset
df = pd.DataFrame({
    'year': [2020, 2020, 2021, 2021, 2022, 2022, 2022],
    'name': ['Alice', 'Bob', 'Alice', 'Bob', 'Alice', 'Charlie', 'Charlie'],
    'ID': ['A001', 'B002', 'A001', 'B002', 'A001', 'C003', 'C003'],
    'value': [10, 15, 20, 25, 30, 35, 40]
})

# Group by 'year' and 'name' columns and calculate the sum of 'value'
df['total'] = df.groupby(['year', 'name', 'ID'])['value'].transform('sum')

# Print the DataFrame
print(df)

Answer 2

如果您需要速度，polars 有一个用 Rust 编写的核心，并且在基准游戏中表现非常好。我个人觉得它的 API 非常直观。这是这个问题的极坐标解决方案，但一如既往，根据您的数据对不同的解决方案进行基准测试。

数据.csv

year,name,id,value
2006,bob,1,3
2006,bob,1,4
2007,sally,1,4

main.py

import polars as pl

q = (
    pl.scan_csv("data.csv")
    .groupby("year", "name", "id")
    .agg((pl.col("value").sum().alias("value_sum"),))
)

q.collect().write_csv("output.csv")

输出.csv

year,name,id,value_sum
2007,sally,1,4
2006,bob,1,7

Answer 3

您可以以 O(n) 的时间复杂度完成此任务。通过使用字典，该解决方案变得微不足道。迭代 csv 中的每一行并提取字段。从那里，使用重复字段作为键关联值。一个简单的例子：

mappings = {}
for line in csv:
    tokens = line.split(",")
    id = tokens[0]
    year = tokens[1]
    value = eval(tokens[2])
    if id in mappings:
         mappings[f"{id}-{year}"] += value
    else:
         mappings[f"{id}-{year}"] = value
print(mappings)

您必须针对您的特定用例进行修改，但向字典添加键的时间复杂度为 O(1)，因此这可能是最简单的“快速”解决方案。如果实施此策略后，您发现需要更快的速度，我建议您研究 C++ 和线程。