I came up with a (perhaps) more pythonic solution which uses groupby
and dropwhile
. This method yields the same result as the below method, but I think it's prettier.. :) Flags, "curr_id" and stuff like that is not very pythonic, and should be avoided if possible!
import csv
from itertools import groupby, dropwhile
goal = 'bob'
ids = {}
with open('my_data.csv') as ifile:
reader = csv.reader(ifile)
for key, rows in groupby(reader, key=lambda r: r[0]):
matched_rows = list(dropwhile(lambda r: r[2] != goal, rows))
if len(matched_rows) > 1:
ids[key] = [row[1] for row in matched_rows[1:]]
print ids
(first solution below)
from collections import defaultdict
import csv
curr_id = None
found = False
goal = 'bob'
ids = defaultdict(list)
with open('my_data.csv') as ifile:
for row in csv.reader(ifile):
if row[0] != curr_id:
found = False
curr_id = row[0]
if found:
ids[curr_id].append(row[1])
elif row[2] == goal:
found = True
print dict(ids)
Output:
{'123': ['2', '8'], '249': ['3']}