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

Question

给定一个文件如下所示:

1440927 1
1727557 3
1440927 2
9917156 4

第一个字段是一个 ID，它在 range(0, 200000000) 内。第二个字段表示一个类型，它是 in range(1, 5)。并且类型 1 和类型 2 属于一个公共(public)类别 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

您在文件上使用迭代器，这意味着您当时只缓冲了几行。每次缓冲区为空时，磁盘都需要寻找并且您的程序必须等待。

200MB 很容易装进你的内存，所以获取所有行会加快速度:

def gen(path):
    # load all the lines, 
    lines = open(path).readlines() 
    split = (line.split() for line in lines)
    return ((int(x), int(y)) for x,y in split)