Extract column from CSV file to use as nodelist in NetworkX
I have a CSV file with two columns: user and locations. I want to create two lists, one with only users and another with only locations, so that I can use the draw_network_nodes (nodelist = ...) function on the networkx to draw users and locations separately as nodes with different shapes and colors (all users will blue boxes and all locations will be red circles).
Also, there is a header in my CSV file, so I don't want the header names to be part of both lists.
+1
source to share
2 answers
Since you have not provided any expected results, I am making some assumptions about them. Let's assume the input file is named data.csv :
user,location
john,seattle
alan,los angeles
trish,new york
The script to split the csv into two files is called 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')
Discussion
- My code is for demo purpose, so it doesn't provide error checking.
- The csv.DictReader () class assumes the first line is the header and uses that as keys for each line
0
source to share
At the top of Hai-Wu's answer:
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
gives:
locations = ('seattle', 'los angeles', 'new york', 'london')
users = ('john', 'alan', 'trish', 'jack')
From:
user,location
john,seattle
alan,los angeles
trish,new york
jack,london
0
source to share