python - 从 CSV 文件中提取列以用作 NetworkX 中的节点列表

Question

我有一个包含 2 列的 CSV 文件：用户和位置。我想创建两个列表：一个只有用户，另一个只有位置，以便我可以使用 networkx 中的 draw_network_nodes(nodelist=...) 函数将用户和位置分别绘制为具有不同形状和颜色的节点（所有用户将是蓝色框，所有位置都将是红色圆圈）。

此外，我的 CSV 文件中有一个标题，因此我不希望标题名称成为任一列表的一部分。

Answer 1

由于您没有提供任何输入、预期输出，因此我对它们进行了一些假设。假设输入文件名为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() 类假定第一行是标题并将其用作每行的键

Answer 2

基于 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