如何提高这个计数程序的性能？

Question

鉴于文件如下所示:

1440927 1
1727557 3
1440927 2
9917156 4

第一个字段是ID in range(0, 200000000).第二个字段代表一种类型,即in range(1, 5).类型1和类型2属于公共类别S1,而类型3和类型4属于公共类别S2.一个ID可能有几个不同类型的记录.该文件大小约为200MB.

问题是计算具有类型1或2的记录的ID的数量,以及具有类型3或4的记录的ID的数量.

我的代码:

def gen(path):
    line_count = 0
    for line in open(path):
        tmp = line.split()
        id = int(tmp[0])
        yield id, int(tmp[1])

max_id = 200000000
S1 = bitarray.bitarray(max_id)
S2 = bitarray.bitarray(max_id)
for id, type in gen(path):
    if type != 3 and type != 4:
        S1[id] = True
    else:
        S2[id] = True

print S1.count(), S2.count()

虽然它给出了答案,但我认为它运行得有点慢.我该怎么做才能让它跑得更快？

编辑: 文件中有重复的记录.我只需要区分S1(类型1和类型2)和S2(类型3和类型4).例如,1440927 1并且1440927 2只计算一次但不计算两次因为它们属于S1.所以我必须存储ID.

Answer 1

如果有足够的内存，您可以dict使用bitarray.bitarray. 它可能会更快：

S1, S2 = {}, {} # dicts are slightly faster than `set()`
with open(path) as f:
     for i, line in enumerate(f, 1):
         id, sep, type = line.partition(" ")
         if type == "1" or type == "2":
            S1[id] = True
         elif type == "3" or type == "4":
            S2[id] = True
         else:
            print "WARNING: unknown type: %r in line %d: %r" % (type, i, line)
print len(S1), len(S2)

或者您可以尝试先对行进行排序：

def gettype(line):
    return line[-1]

S1, S2 = 0, 0
with open(path) as f:
     lines = f.read().splitlines()

lines.sort(key=gettype)
for type, group in itertools.groupby(lines, gettype):
    ids = (line.partition(" ")[0] for line in group)
    if type == "1" or type == "2":
       S1 += len(set(ids))
    elif type == "3" or type == "4":
       S2 += len(set(ids))
    else:
       assert 0, (type, list(ids))

print S1, S2

第二种方法的渐近复杂度更差。

您可以使用line_profiler找出瓶颈所在。