如何多次按索引获取Python生成器中的指定位置元素

问题原因 #

生成器是单向惰性迭代的结构，所有元素只会生成一次，迭代过的元素不会被缓存，迭代指针也无法回退。
你第一次调用next(itertools.islice(x, 50, None))时，islice会直接消耗掉生成器的前50个元素（索引0~49），再取出索引50的元素返回，此时生成器的指针已经停在索引51的位置。
后续你再调用next(itertools.islice(x, 25, None))时，islice是从当前指针位置（索引51）往后再跳过25个元素，取到的就是50 + 1 + 25 = 76，自然不是你预期的索引25对应的元素。

解决方案 #

根据你的使用场景可以选以下三种方案：

• 方案1：生成器元素量不大的情况下直接转为列表，是最简单的随机访问方案
  列表支持O(1)复杂度的下标随机访问，直接转列表后按需取索引即可：

  x = (i for i in range(100))
  lst = list(x)
  # 直接按索引取值
  print(lst[50]) # 输出50
  print(lst[25]) # 输出25
  

• 方案2：元素量较大不想一次性全加载的话，用itertools.tee创建独立迭代器
  你需要取几次不同位置的元素，就提前复制几个独立的迭代器，每个迭代器的指针互不影响：

  import itertools
  x = (i for i in range(100))
  # 复制2个独立迭代器
  iter1, iter2 = itertools.tee(x, 2)
  # 用第一个迭代器取索引50的元素
  print(next(itertools.islice(iter1, 50, None))) # 输出50
  # 用第二个迭代器取索引25的元素
  print(next(itertools.islice(iter2, 25, None))) # 输出25
  

  注意：如果多个迭代器的进度差非常大，进度靠前的迭代器走过的元素会被临时缓存，仍会占用一定内存。

• 方案3：需要频繁按任意索引取值，可自行实现带缓存的生成器包装类
  把已经生成过的元素存在缓存列表里，取索引时优先读缓存，没有再继续迭代生成器：

  class CachedGenerator:
      def __init__(self, gen):
          self.gen = gen
          self.cache = []
      
      def __getitem__(self, idx):
          # 迭代生成器直到缓存覆盖目标索引
          while len(self.cache) <= idx:
              self.cache.append(next(self.gen))
          return self.cache[idx]
  

  使用方式如下：

  x = CachedGenerator(i for i in range(100))
  print(x[50]) # 输出50
  print(x[25]) # 输出25
  

内容的提问来源于stack exchange，提问作者sho ryu