打印匹配标题值的列

Question

我有两个csv文件。

文件1:

id,site,longitude,latitude             
**9936**,north,18.2,62.8              
5856,north,17.4914,63.0167             
**1298**,north,18.177,62.877   

文件2:

chr,loc,4678,**1298**,2295,**9936**,7354             
chr1,849,0,0,0,0,0,             
chr1,3481,1,1,0,1,1                             
chr1,3491,0,2,0,2,0,             

我想把文件1中第一列的id和文件2中的行进行匹配（用**标记的部分），如果匹配成功，就打印出这一行和相关的内容。

输出结果:

chr,loc,**1298**,**9936**            
chr1,849,0,0             
chr1,3481,1,1                             
chr1,3491,0,2

我在尝试用python来实现这个功能。

import csv

f1 = file('inFile.csv', 'rb')                 
f2 = file('inFile2.csv', 'rb')               
f3 = file('outFile.csv', 'wb')                           
c1 = csv.reader(f1)            
c2 = csv.reader(f2)                 
c3 = csv.writer(f3)              

matched_rows = [ row for row in c2 if row[2:6] in c1]           
for row in matched_rows:                                                  
    c3writerow[matched_rows]

但是不幸的是，它没有成功。

Answer 1

你需要先从文件1中加载那一列数据，并把它存储成一种方便查找的格式。这里可以用一个叫做set的东西：

with open('inFile.csv', 'rb') as ids_file:
    reader = csv.reader(ids_file)
    next(reader, None)  # skip the first row
    ids = {r[0] for r in reader}

现在你可以测试一下你的匹配列了：

from operator import itemgetter

with open('inFile2.csv', 'rb') as f2, file('outFile.csv', 'wb') as outf:
    reader = csv.reader(f2)
    writer = csv.writer(outf)

    headers = next(reader, [])
    # produce indices for what headers are present in the ids set
    matching_indices = [i for i, header in enumerate(headers[2:], 2) if header in ids]
    selector = itemgetter(0, 1, *matching_indices)
    # write selected columns to output file
    writer.writerow(selector(headers))
    writer.writerows(selector(row) for row in reader)

在你的示例数据上演示一下：

首先，生成第一列的一个集合：

>>> ids_file = '''\
... id,site,longitude,latitude
... 9936,north,18.2,62.8
... 5856,north,17.4914,63.0167
... 1298,north,18.177,62.877
... '''.splitlines()
>>> reader = csv.reader(ids_file)
>>> next(reader, None)
['id', 'site', 'longitude', 'latitude']
>>> ids = {r[0] for r in reader}
>>> ids
set(['5856', '9936', '1298'])

然后用这些数据来生成一个选择器，使用operator.itemgetter()：

>>> from operator import itemgetter
>>> f2 = '''\
... chr,loc,4678,1298,2295,9936,7354
... chr1,849,0,0,0,0,0,
... chr1,3481,1,1,0,1,1
... chr1,3491,0,2,0,2,0,
... '''.splitlines()
>>> reader = csv.reader(f2)
>>> headers = next(reader, [])
>>> matching_indices = [i for i, header in enumerate(headers[2:], 2) if header in ids]
>>> matching_indices
[3, 5]
>>> selector = itemgetter(0, 1, *matching_indices)

现在你可以用这个对象来选择你想要的列，然后写入输出的CSV文件中：

>>> selector(headers)
('chr', 'loc', '1298', '9936')
>>> selector(next(reader))
('chr1', '849', '0', '0')
>>> selector(next(reader))
('chr1', '3481', '1', '1')
>>> selector(next(reader))
('chr1', '3491', '2', '2')