如何更高效地在Python中对CSV文件的列求和

Question

这是我的数据：

Year    Country     Albania     Andorra     Armenia     Austria   Azerbaijan
2009    Lithuania      0           0           0           0           1    
2009    Israel         0           7           0           0           0    
2008    Israel         1           2           2           0           4
2008    Lithuania      1           5           1           0           8    

其实，这是一份csv文件，分隔符是逗号，所以原始数据是：

Year,Country,Albania,Andorra,Armenia,Austria,Azerbaijan
2009,Lithuania,0,0,0,0,1
2009,Israel,0,7,0,0,0
2008,Israel,1,2,2,0,4
2008,Lithuania,1,5,1,0,8

这里列表的第一个元素表示立陶宛的列总和，第二个元素表示以色列的列总和（针对阿尔巴尼亚这一列）？

我刚开始学python，对很多技巧还不太了解。我知道我可能在代码中把事情搞得太复杂了。

我想得到这个：

final_dict = {Albania: [1, 1], Andorra: [5, 9], Armenia: [1, 2], Austria: [0, 0], Azerbaijan: [9, 4]}

输出解释：对于第一行的每个国家（阿尔巴尼亚、安道尔、亚美尼亚、奥地利和阿塞拜疆），我想要从国家列中得到各国的总和。

Andorra: [5,9] 
# 5 is sum for Lithuania in Andorra column
# 9 is sum for Israel in Andorra column

Answer 1

你可以使用 collections模块中的defaultdict，请在StackOverflow上搜索

python defaultdict

你会找到很多有用的例子，这里是我的回答

import csv
from collections import defaultdict

# slurp the data
data = list(csv.reader(open('points.csv')))

# massage the data
for i, row in enumerate(data[1:],1):
    data[i] = [int(elt) if elt.isdigit() else elt for elt in row]

points = {} # an empty dictionary
for i, country in enumerate(data[0][2:],2):
    # for each country, a couple country:defaultdict is put in points
    points[country] = defaultdict(int)
    for row in data[1:]:
        opponent = row[1]
        points[country][opponent] += row[i]

# here you can  post-process  points as you like,
# I'll simply print out the stuff
for country in points:
    for opponent in points[country]:
        print country, "vs", opponent, "scored",
        print points[country][opponent], "points."

你数据的示例输出是

Andorra vs Israel scored 9 points.
Andorra vs Lithuania scored 5 points.
Austria vs Israel scored 0 points.
Austria vs Lithuania scored 0 points.
Albania vs Israel scored 1 points.
Albania vs Lithuania scored 1 points.
Azerbaijan vs Israel scored 4 points.
Azerbaijan vs Lithuania scored 9 points.
Armenia vs Israel scored 2 points.
Armenia vs Lithuania scored 1 points.

编辑

如果你不想用 defaultdict，你可以使用普通 dict 的 .get 方法，这样如果 key:value 对没有初始化，你可以得到一个可选的默认值。

    points[country] = {} # a standard empty dict
    for row in data[1:]:
        opponent = row[1]
        points[country][opponent] = points[country].get(opponent,0) + row[i]

如你所见，这种方法稍微复杂一点，但还是可以处理的。

Answer 2

你可以使用Pandas模块，这个模块非常适合这种类型的应用：

import pandas as pd

df = pd.read_csv('songfestival.csv')
gb = df.groupby('Country')
res = pd.concat([i[1].sum(numeric_only=True) for i in gb], axis=1).T
res.pop('Year')
order = [i[0] for i in gb]

print(order)
print(res)

#['Israel', 'Lithuania']
#   Albania  Andorra  Armenia  Austria  Azerbaijan
#0        1        9        2        0           4
#1        1        5        1        0           9

要查询每一列的结果，你只需要这样做：

print(res.Albania)
print(res.Andorra)
...

Answer 3

好的，你想要按年份把数据合并在一起：

import csv
from collections import defaultdict

with open("songfestival.csv", "r") as ifile:
    reader = csv.DictReader(ifile)
    country_columns = [k for k in reader.fieldnames if k not in ["Year","Country"]]
    data = defaultdict(lambda:defaultdict(int))
    for line in reader:
        curr_country = data[line["Country"]]
        for country_column in country_columns:
            curr_country[country_column] += int(line[country_column])

    with open("songfestival_aggr.csv", "w") as ofile:
        writer = csv.DictWriter(ofile, fieldnames=country_columns+["Country"])
        writer.writeheader()
        for k, v in data.items():
            row = dict(v)
            row["Country"] = k
            writer.writerow(row)

我就随便把结果输出到另一个csv文件里了。你的数据结构很容易出错，因为它依赖于列的顺序。最好使用一个字典里的字典来给合并的结果命名——可以参考@gboffi在你问题下的评论。