Python - 从csv中条件求和

Question

我有一个CSV文件，其中有一列叫做“cost”（成本），我想根据另一列“factory”（工厂）来计算总成本，也就是想要按工厂来分解成本。我的数据行大概是这样的，每个工厂都有多个成本：

Factory,Cost,Cost_Type
Bali,23,0
Sydney,21,1
Sydney,4,2
Denver,8,1
Bali,9,1

我希望能快速计算每个工厂的总成本，并把这些值保存到一个变量里。我觉得可以通过先列出工厂的列表，然后再遍历CSV文件来实现。现在我已经做到这个步骤：

factories= ['Bali', 'Sydney', 'Denver']
totalcost = 0
balicost = 0
sydneycost = 0
denvercost = 0

for factory in factories:
    for row in csv.reader(costcsv):
        if row[0] == factory:

我现在遇到的问题是，我不知道怎么改变用于不同工厂的变量，比如balicost、sydneycost和denvercost。我之前做的一个简化版本，只是计算了成本列的总和，代码如下：

for row in csv.reader(costcsv):
        totalcost += float(row[1])

我也很欢迎其他的方法（我觉得字典可能会有用），如果能给我一些指引，我会很感激。

Answer 1

你的csv文件应该是这样的：

Factory,Cost
Bali,23
Sydney,21
Sydney,4
Denver,8
Bali,9

在Python中，你可以这样做：

import csv

factories= ['Bali', 'Sydney', 'Denver']
totalcost = 0

sums = {}

with open('file.csv', 'rb') as f:
    f.next()                        # Jump to second row -> first : header
    reader = csv.reader(f)
    for row in reader:
        if row[0] not in sums:
            sums[row[0]] = int(row[1])
        else:
            sums[row[0]] += int(row[1])


for key,value in sums.items():
    totalcost = totalcost  + int(value)

结果看起来是这样的：

print sums
>{'Denver': 8, 'Sydney': 25, 'Bali': 32}
print totalcost
>65

Answer 2

最简单的方法是用一个字典来记录每个工厂的数量：

factoriescost = {}
for row in cvs.reader(costcsv):
    factory = row[0]
    if factory not in ('Bali', 'Sydney', 'Denver'):
        continue
    factorycost = factoriescost.get(factory, 0)
    factoriescost[factory] = factorycost + float(row[1])
totalcost = sum(factoriescost.itervalues())

然后你可以用 factoriescost 来获取某个特定工厂的总费用：

>>> print totalcost, factoriescost
65.0 {'Denver': 8.0, 'Sydney': 25.0, 'Bali': 32.0}
>>> print factoriescost['Bali']
32.0

Answer 3

与其使用多个单独的变量，不如考虑使用字典，或者更简单的，使用collections.defaultdict：

from collections import defaultdict

costs = defaultdict(float)

for line in csv.reader(costcsv):
    if len(line) == 2:
        factory, costs = line
        costs[factory] += float(cost)

这样你就可以得到一个输出，能够选择任何工厂（不仅仅是你现在硬编码的那三个），并计算出总费用。

cost["denver"] == 8.0

Answer 4

你可以像下面这样使用字典。这个代码使用了一个 try 循环来计算字典中工厂的成本，如果工厂不在字典里，就会出现 KeyError 错误，这时候就直接把这个工厂添加进去。

a = [['Bali', 23],
     ['Sydney', 21],
     ['Sydney', 4],
     ['Denver', 8],
     ['Bali', 9]]

factories = dict()

for factory, cost in a:
    try:
        factories[factory] += cost
    except KeyError:
        factories[factory] = cost

print(factories)
# {'Denver': 8, 'Sydney': 25, 'Bali': 32}

在你的例子中，你可以用一个合适的循环来替换掉 for 循环，使用 csv.reader()，大致可以这样写：

for factory, cost in csv.reader(costcsv):
    try:
        ...

Answer 5

[社区维基，因为这个话题稍微偏离了主题。]

当你在Python中处理表格数据时，建议你使用pandas这个库。你想要做的操作是对数据进行分组求和，这个操作可以用两行代码轻松完成：

df = pd.read_csv("factories.csv")
by_factory = df.groupby("Factory")["Cost"].sum()

这样会生成一个Series对象，你可以像查字典一样来访问它：

>>> by_factory
Factory
Bali       32
Denver      8
Sydney     25
Name: Cost, dtype: int64
>>> by_factory["Bali"]
32

---

更新一下，使用更新后的数据——如果你还想处理Cost_Type，你有几种选择。一个方法是只选择Cost_Type等于1的行：

>>> df[df.Cost_Type == 1]
  Factory  Cost  Cost_Type
1  Sydney    21          1
3  Denver     8          1
4    Bali     9          1

[3 rows x 3 columns]
>>> df[df.Cost_Type == 1].groupby("Factory")["Cost"].sum()
Factory
Bali        9
Denver      8
Sydney     21
Name: Cost, dtype: int64

或者你可以扩展分组操作，同时对Factory和Cost_Type进行分组：

>>> df.groupby(["Cost_Type", "Factory"])["Cost"].sum()
Cost_Type  Factory
0          Bali       23
1          Bali        9
           Denver      8
           Sydney     21
2          Sydney      4
Name: Cost, dtype: int64