从CSV文件提取列作为NetworkX中的节点列表

Question

我有一个CSV文件，里面有两列：用户和地点。我想创建两个列表：一个只包含用户，另一个只包含地点。这样我就可以使用networkx中的draw_network_nodes(nodelist=...)函数，把用户和地点分别画成不同形状和颜色的节点（所有用户都是蓝色方块，所有地点都是红色圆圈）。

另外，我的CSV文件里有一个标题行，所以我不想把标题的名字包含在这两个列表里。

Answer 1

在Hai Vu的回答基础上：

import csv
def reader(filename):
    for (lineno, line) in enumerate(open(filename)):
        if lineno > 0: # skip header
            yield line

filename = "locations.csv"
(users, locations) = zip(*( row for row in csv.reader(reader(filename))))
print "users     =", users
print "locations =", locations

结果是：

locations = ('seattle', 'los angeles', 'new york', 'london')
users     = ('john', 'alan', 'trish', 'jack')

来自：

user,location
john,seattle
alan,los angeles
trish,new york
jack,london

Answer 2

因为你没有提供输入和预期的输出，我只能做一些假设。假设输入文件叫data.csv：

user,location
john,seattle
alan,los angeles
trish,new york

把csv文件分成两个文件的脚本叫csv_split.py：

import csv

with open('data.csv') as csv_in,          \
     open('users.txt', 'w') as users_out, \
     open('locations.txt', 'w') as locations_out:
    csv_dict_reader = csv.DictReader(csv_in)
    for line in csv_dict_reader:
        users_out.write(line['user'] + '\n')
        locations_out.write(line['location'] + '\n')

讨论

• 我的代码只是用来演示的，所以没有做任何错误检查。
• csv.DictReader()这个类假设第一行是标题，并用它作为每一行的键。