Извлечь столбец из файла CSV для использования в качестве нодлиста в NetworkX

У меня есть файл CSV с двумя столбцами: пользователь и местоположение. Я хочу создать два списка: один только с пользователями, а другой только с местоположениями, чтобы я мог использовать функцию draw_network_nodes(nodelist=...) в networkx для рисования пользователей и местоположений отдельно в виде узлов с разными формами и цветами (все пользователи будут синими прямоугольниками, а все места будут красными кружками).

Кроме того, в моем файле CSV есть заголовок, поэтому я не хочу, чтобы имена заголовков были частью каких-либо списков.


person user2096629    schedule 21.02.2013    source источник
comment
Пожалуйста, отредактируйте свой пост, чтобы включить вашу попытку решения. Спасибо.   -  person mechanical_meat    schedule 21.02.2013
comment
docs.python.org/2/library/csv.html Вы приняли посмотреть на это? Здесь есть несколько довольно ярких примеров.   -  person Daan Lubbers    schedule 21.02.2013


Ответы (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() предполагает, что первая строка является заголовком, и использует ее в качестве ключей для каждой строки.
person Hai Vu    schedule 21.02.2013

Опираясь на ответ Хай Ву:

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
person sotapme    schedule 21.02.2013