在Python中快速迭代可迭代对象的前n个项目

Question

我在寻找一种在可迭代对象中快速遍历前 n 个项目的“Pythonic”方法（更新：通常情况下不是列表，因为列表的处理比较简单），而且尽可能快地做到这一点对我来说非常重要。现在我用的方法是：

count = 0
for item in iterable:
 do_something(item)
 count += 1
 if count >= n: break

我觉得这个方法看起来不太整洁。另一种做法是：

for item in itertools.islice(iterable, n):
    do_something(item)

这个方法看起来不错，但问题是它在使用某些生成器时速度够快吗？比如说：

pair_generator = lambda iterable: itertools.izip(*[iter(iterable)]*2)
for item in itertools.islice(pair_generator(iterable), n):
 so_something(item)

与第一种方法相比，这个方法运行得快吗？ 有没有更简单的方法可以做到这一点？

Answer 1

你可以使用 enumerate 这个功能，来写一个和你现在的循环基本一样的代码，但会更简单、更符合Python的风格：

for idx, val in enumerate(iterableobj):
    if idx > n:
        break
    do_something(val)

Answer 2

itertools 通常是最快的解决方案，当它可以直接使用的时候。

显然，唯一的方法是进行基准测试——比如，把代码保存到 aaa.py 文件里。

import itertools

def doit1(iterable, n, do_something=lambda x: None):
  count = 0
  for item in iterable:
   do_something(item)
   count += 1
   if count >= n: break

def doit2(iterable, n, do_something=lambda x: None):
  for item in itertools.islice(iterable, n):
      do_something(item)

pair_generator = lambda iterable: itertools.izip(*[iter(iterable)]*2)

def dd1(itrbl=range(44)): doit1(itrbl, 23)
def dd2(itrbl=range(44)): doit2(itrbl, 23)

然后看看结果……：

$ python -mtimeit -s'import aaa' 'aaa.dd1()'
100000 loops, best of 3: 8.82 usec per loop
$ python -mtimeit -s'import aaa' 'aaa.dd2()'
100000 loops, best of 3: 6.33 usec per loop

所以很明显，在这里 itertools 更快——用你自己的数据进行基准测试来验证一下。

顺便说一下，我发现从命令行使用 timeit 要方便得多，所以我总是这样用——它会根据你想测量的速度，运行适当数量的循环，比如10次、100次、1000次等等——在这里，为了区分一毫秒和半毫秒的差别，跑十万次循环差不多就够了。

Answer 3

for item in itertools.islice(iterable, n): 是一种最简单、最直接的方法来实现这个功能。它适用于任何可迭代的对象，并且时间复杂度是 O(n)，这在合理的解决方案中是正常的。

当然，可能会有其他方法性能更好，但我们需要通过计时来确认。如果你没有对代码进行性能分析，发现这个调用是个瓶颈，我不建议你去计时。除非这个调用在一个内层循环里，否则很难说它会成为瓶颈。过早优化是万恶之源。

---

如果我真的要寻找其他解决方案，我会考虑像 for count, item in enumerate(iterable): if count > n: break ... 和 for i in xrange(n): item = next(iterator) ... 这样的方式。我不敢保证这些方法会有帮助，但如果我们真的想比较一下，试试也无妨。如果我发现这个调用在内层循环中确实是个瓶颈（这真的是你的情况吗？），我还会尝试把 islice 从全局的 itertools 中获取的过程简化，直接将这个函数绑定到一个局部变量上。

这些都是在你确认它们会有帮助之后才去做的事情。很多人会在其他时候尝试这些方法，但这并不会让他们的程序明显更快，反而可能让程序变得更糟。