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

【问题标题】：How to improve performance of this counting program?如何提高这个计数程序的性能？
【发布时间】：2012-01-13 22:16:54
【问题描述】：

给定一个文件如下所示：

第一个字段是一个 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。

【问题讨论】：

您可以使用分析器。您可以删除 id=int( ... 并改用 yield int(tmp[0], ...。您可以使用if type <= 2 而不是两个比较。您可以完全删除生成器并将代码内联到 with open( ... ) as f: 块中。试一试。下面的评论也有一个好点，关于 bitarray ^^
您有什么理由使用位数组来标记索引吗？否则，您可以简单地增加一个计数器，而不是将条目设置为“True”。这应该可以提高性能。
+1 使用分析器。瓶颈在哪里？是S1和S2的分配吗？另外，请考虑以下问题：（几乎）0-200000000 中的所有数字都存在吗？如果不是，请考虑另一种数据类型。每个id可以出现多次吗？如果没有，请考虑完全放弃阵列并仅使用计数器。或者也许这是一个您已经有了最佳解决方案的问题。对于非常大的文件，您的瓶颈很可能是磁盘 I/O，这将需要您购买更好的磁盘进行优化。
@Boris 我必须存储 ID，因为有重复的记录。例如，在文件样本 1440927 中，应该只计算一次，而不是两次。因为类型 1 和类型 2 都属于 S1。

标签： python performance bitarray

【解决方案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 找出您的瓶颈所在。

【讨论】：

【解决方案2】：

您正在对文件使用迭代器，这意味着您一次只缓冲几行。每次缓冲区为空时，磁盘都需要查找，您的程序必须等待。

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)

【讨论】：

看起来您在解决方案中使用了 600MB。
@hochl：好的，我将列表理解更改为生成器表达式。现在它应该使用 200MB 来存储lines。
你不能确定哪个更快for line in f.readlines() 或for line in f 除非分析器说。文件迭代器使用READAHEAD_BUFSIZE (8192) 这意味着在这种情况下一次数百行。

【解决方案3】：

你是否依赖 Python？

egrep -e "[12]$" filename.txt | cut -d " " -f 1 | sort -u | wc -l

egrep -e "[34]$" filename.txt | cut -d " " -f 1 | sort -u | wc -l

这两个命令计算文件名.txt 中每行末尾出现 ("1" 或 "2") 和 ("3" 或 "4") 的次数，同时忽略重复的第一个字段。

可能比 Python 快……

【讨论】：

uniq 需要排序输入，而 OP 没有。您可以在管道中添加sort...
你被 Python 绑定了吗？ vs. 你绑定到 Linux 了吗？ :)
@warvariuc：我的 Windows 桌面在命令行上有一个 grep -E 可用......你的意思是什么？
@MattH，我的意思是：什么更好 - 绑定到单独的程序，还是用 Python 做所有事情？
@warvariuc：我倾向于说适合工作的正确工具是最好的方法。