Python中的Vlookup

Question

我刚开始学Python，想尽快掌握它。我知道怎么用bash来解决我的问题，现在想用Python来实现。

我有一个数据文件（在这个例子中是data_array.csv）和一个索引文件index.csv。我想从数据文件中提取那些在索引文件中有相同ID的数据，并把它们存到一个新文件Out.txt里。如果某些ID在数据文件中没有值，我还想在Out.txt里填上NA。我知道怎么处理一列数据，但我的数据有超过1000列（从1到1344列）。我希望你能帮我写一个脚本，让这个过程更快。我的数据文件、索引ID和想要的输出如下。

data_array.csv

Id  1   2   3   .   .   1344
1   10  20  30  .   .   -1
2   20  30  40  .   .   -2
3   30  40  50  .   .   -3
4   40  50  60  .   .   -4
6   60  60  70  .   .   -5
8   80  70  80  .   .   -6
10  100 80  90  .   .   -7

index.csv

Id
1
2
8
9
10

所需的输出是

Out.txt

Id  1   2   3   .   .   1344
1   10  20  30  .   .   -1
2   20  30  40  .   .   -2
8   80  70  80  .   .   -6
9   NA  NA  NA          NA
10  100 80  90  .   .   -7

我试过

#! /usr/bin/python

import csv

with open('data_array.csv','r') as lookuplist:
    with open('index.csv', "r") as csvinput:
        with open('VlookupOut','w') as output:

            reader = csv.reader(lookuplist)
            reader2 = csv.reader(csvinput)
            writer = csv.writer(output)

            for i in reader2:
                for xl in reader:
                    if i[0] == xl[0]:
                        i.append(xl[1:])
                        writer.writerow(i)

但它只处理了第一行。我希望程序能处理我数据文件的所有行和列。

Answer 1

当你用 for xl in reader 来读取一个CSV文件时，它会一行一行地读取，直到文件的最后一行。但这个过程只会进行一次。如果你想让它重新回到CSV文件的第一行，可以使用 .seek(0) 这个命令。

#! /usr/bin/python

import csv

with open('data_array.csv','r') as lookuplist:
    with open('index.csv', "r") as csvinput:
        with open('VlookupOut','w') as output:

            reader = csv.reader(lookuplist)
            reader2 = csv.reader(csvinput)
            writer = csv.writer(output)

            for i in reader2:
                for xl in reader:
                    if i[0] == xl[0]:
                        i.append(xl[1:])
                        writer.writerow(i)
                lookuplist.seek(0)

Answer 2

它只输出了第一行，因为在第一次执行 xl in reader 后，你已经到了文件的末尾。之后你需要把指针移回文件的开头。为了提高效率，你可以先把 csvinput 读入一个字典，然后通过字典查找你需要的行：

#! /usr/bin/python

import csv

with open('data_array.csv','r') as lookuplist:
    with open('index.csv', "r") as csvinput:
        with open('VlookupOut','w') as output:

            reader = csv.reader(lookuplist)
            reader2 = csv.reader(csvinput)
            writer = csv.writer(output)

            d = {}
            for xl in reader2:
                d[xl[0]] = xl[1:]

            for i in reader:
                if i[0] in d:
                    i.append(d[i[0]])
                writer.writerow(i)