需要更高效的方式解析Python中的CSV文件

Question

这是一个示例的csv文件

id, serial_no
2, 500
2, 501
2, 502
3, 600
3, 601

我想要的输出结果是这样的（在一个id列表中列出serial_no）：

[2, [500,501,502]]
[3, [600, 601]]

我已经实现了我的解决方案，但代码太多了，我相信还有更好的方法。我还在学习Python，很多技巧还不太懂。

file = 'test.csv'

data = csv.reader(open(file))
fields = data.next()

for row in data:
  each_row = []     
    each_row.append(row[0])
    each_row.append(row[1])
    zipped_data.append(each_row)
for rec in zipped_data:
  if rec[0] not in ids:
    ids.append(rec[0])
for id in ids:
    for rec in zipped_data:
      if rec[0] == id:
        ser_no.append(rec[1])
  tmp.append(id)
  tmp.append(ser_no)
  print tmp
  tmp = []
  ser_no = []

**为了简化代码，我省略了变量初始化的部分

print tmp

这段代码给了我上面提到的输出结果。我知道还有更好的方法，或者说更符合Python风格的方法。现在的代码看起来太乱了！如果有建议就太好了！

Answer 1

这是我写的一个版本，不过看起来已经有很多人回答这个问题了。

你可以试试使用csv.DictReader，这样可以通过字段名称（也就是表头/第一行）轻松访问每一列的数据。

#!/usr/bin/python
import csv

myFile = open('sample.csv','rb')
csvFile = csv.DictReader(myFile)
# first row will be used for field names (by default)

myData = {}

for myRow in csvFile:
    myId = myRow['id']
    if not myData.has_key(myId): myData[myId] = []
    myData[myId].append(myRow['serial_no'])

for myId in sorted(myData):
    print '%s %s' % (myId, myData[myId])

myFile.close()

Answer 2

与其用一个普通的列表，我会用一个叫做 collections.defaultdict(list) 的东西。这样的话，你可以直接在这个值上使用 append() 方法来添加内容。

result = collections.defaultdict(list)
for row in data:
  result[row[0]].append(row[1])

Answer 3

from collections import defaultdict

records = defaultdict(list)

file = 'test.csv'

data = csv.reader(open(file))
fields = data.next()

for row in data:
    records[row[0]].append(row[1])

#sorting by ids since keys don't maintain order
results = sorted(records.items(), key=lambda x: x[0])
print results

如果你希望序列号的列表是唯一的，只需要把 defaultdict(list) 替换成 defaultdict(set)，然后把 records[row[0]].append(row[1]) 替换成 records[row[0]].add(row[1])。