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

Question

给定一个看起来像这样的文件：

1440927 1
1727557 3
1440927 2
9917156 4

第一个字段是ID，其范围为0至200000000。第二个字段表示类型，其范围为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。

- amazingjxq

---

你可以使用性能分析器。你可以删除id=int(...，并改为使用yield int(tmp[0], ...。你可以使用if type <= 2代替两个比较。你可以完全删除生成器，并将代码内联到with open(...) as f:块中。试一试吧。下面的评论也有一个关于bitarray的好观点^^ - hochl

你使用 bitarray 标记索引有什么原因吗？否则，你可以简单地增加一个计数器而不是将条目设置为“True”。这应该能提高性能。 - Boris

在使用分析器时加1。瓶颈在哪里？是S1和S2的分配吗？此外，请考虑以下问题：0-200000000中是否（几乎）所有数字都存在？如果不是，请考虑另一种数据类型。每个ID是否可以出现多次？如果不是，请考虑完全放弃数组，只使用计数器。或者也许这是一个您已经拥有最佳解决方案的问题。对于非常大的文件，您的瓶颈可能是磁盘I / O，这将需要您购买更好的磁盘以进行优化。 - Emil Vikström

@Boris，我必须存储这些ID，因为存在重复记录。例如，在样本文件中，1440927应该只计算一次，而不是两次。因为类型1和2都属于S1。 - amazingjxq

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)

- Jochen Ritzel

---

你的解决方案似乎使用了600MB。 - hochl

@hochl: 好的，我把列表推导式改成了生成器表达式。现在它应该使用200MB存储“lines”。 - Jochen Ritzel

除非使用分析器，否则无法确定哪个更快for line in f.readlines()或for line in f。文件迭代器使用READAHEAD_BUFSIZE（8192），这意味着在这种情况下每次读取数百行。 - jfs

Answer 2

你是否局限于使用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

这两个命令可以统计你的文件filename.txt每一行结尾处出现("1"或"2")并且出现("3"或"4")的次数，同时忽略重复的第一个字段。可能比Python更快...

- eumiro

---

“uniq”需要经过排序的输入，而原帖中并没有排序。您可以在管道中添加一个“sort”... - Karl Knechtel

你是否被绑定在Python上？vs. 你是否被绑定在Linux上？ :) - warvariuc

@warvariuc：我的Windows桌面可以在命令行上使用grep -E...你的意思是什么？ - MattH

@MattH，我的观点是：什么更好-绑定到一个单独的程序还是在Python中完成所有操作？ - warvariuc

@warvariuc：我倾向于认为，选择合适的工具来完成相应的工作是最佳方法。 - MattH

Answer 3

如果有足够的内存，你可以使用 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来找出瓶颈在哪里。

- jfs