如何用Python高效找到两个大文件的交集？

Question

我有两个很大的文件。它们的内容看起来像这样：

134430513
125296589
151963957
125296589

这些文件里有一堆乱七八糟的ID。有些ID在同一个文件里可能会出现多次。

现在我想找出这两个文件的交集部分，也就是在两个文件中都出现的ID。

我把这两个文件的内容分别读入了两个集合，叫做s1和s2。然后通过s1.intersection(s2)来获取交集。但是这样做占用了很多内存，而且似乎还挺慢的。

那么有没有更好的、符合Python风格的方法来实现这个呢？如果文件里的ID多到无法在有限的内存中读入一个set，我该怎么办呢？

补充：我用生成器把文件读入了两个集合：

def id_gen(path):
    for line in open(path):
        tmp = line.split()
        yield int(tmp[0])

c1 = id_gen(path)
s1 = set(c1)

所有的ID都是数字。而且最大的ID可能是5000000000。如果使用位数组（bitarray），那会占用更多的内存。

Answer 1

这个算法其实并不是专门针对Python的，而是一个通用的方法，适合在内存中无法一次性存下所有ID的情况。如果整数的范围有限，可以使用一个大的位数组（bitarray）。首先，你读取第一个文件，并在这个bitarray中标记出这些整数是存在的。接着，你再读取第二个文件，输出所有在bitarray中也存在的数字。

如果这样还不够，你可以通过多次扫描来分割范围。比如在第一次扫描时，只考虑小于0x200000000的整数（这时需要一个1GB的bitarray）。然后你重置bitarray，再次读取文件，这次只考虑从0x200000000到0x400000000之间的整数（处理这些整数时要先减去0x200000000）。

这样一来，你就可以处理大量的数据，而且运行时间也比较合理。

一个单次扫描的示例代码如下：

import bitarray
r = bitarray.bitarray(5000000000)

for line in open(file1):
    r[int(line)] = True

for line in open(file2):
    if r[int(line)]:
        print line

Answer 2

set(open(file1)) & set(open(file2))

使用 intersection 是最符合Python风格的方法。你可以通过以下方式来加快速度：

set(int(x) for x in open(file1)) & set(int(x) for x in open(file2))

因为这样你就可以存储和比较整数，而不是字符串。当然，这只有在所有的id都是数字的情况下才有效。

如果速度还是不够快，你可以尝试一种稍微不同的写法：

# heuristic: set smaller_file and larger_file by checking the file size
a = set(int(x) for x in open(smaller_file))
# note: we're storing strings in r
r = set(x for x in open(larger_file) if int(x) in a)

如果两个文件都保证没有重复项，你也可以使用列表来加速处理：

a = set(int(x) for x in open(smaller_file))
r = [x for x in open(larger_file) if int(x) in a]

一定要测试不同的解决方案，看看你是不是在等硬盘或网络的输入。

Answer 3

其他人已经展示了在Python中更常见的做法，但如果数据的大小确实太大，你可以使用系统工具来排序和去重，然后利用文件是一个迭代器的特点，它一次返回一行，做类似这样的操作：

import os
os.system('sort -u -n s1.num > s1.ns')
os.system('sort -u -n s2.num > s2.ns')
i1 = open('s1.ns', 'r')
i2 = open('s2.ns', 'r')
try:
    d1 = i1.next()
    d2 = i2.next()
    while True:
        if (d1 < d2):
            d1 = i1.next()
        elif (d2 < d1):
            d2 = i2.next()
        else:
            print d1,
            d1 = i1.next()
            d2 = i2.next()
except StopIteration:
    pass

这样做可以避免在内存中同时加载多于一行的数据（对于每个文件来说），而且系统排序的速度应该比Python的任何方法都快，因为它是专门为这个任务优化的。