在 Python 中获取生成器的第 N 个元素

Python 1年前 19

获取生成器的第 N 个元素：

使用 islice() 函数获取从索引 N 开始的迭代器。
将迭代器传递给 next() 函数。
next() 函数将返回生成器的第 N 个元素。

from itertools import islice


def g():
    yield 1
    yield 2
    yield 3


gen = g()

second = next(
    islice(gen, 1, None),
    None
)
print(second)  # 👉️ 2

itertools.islice() 方法采用迭代器和可选的开始和停止索引。

该方法创建了一个迭代器，该迭代器从可迭代对象中返回选定的元素。

我们使用了 0 的开始索引和 None 的停止索引。

当 stop 为 None 时，迭代将继续，直到迭代器耗尽。

对于无限长的生成器，迭代会在指定位置停止。

next() 函数从提供的迭代器中返回下一项。

该函数可以传递一个默认值作为第二个参数。

如果迭代器耗尽或为空，则返回默认值。

如果迭代器已用尽或为空且未提供默认值，则会引发 StopIteration 异常。

我们使用 None 作为默认值，但我们可以使用适合自己用例的任何其他值。

from itertools import islice


def g():
    yield 1
    yield 2
    yield 3


gen = g()

item = next(
    islice(gen, 100, None),
    'default value'
)
print(item)  # 👉️ default value

如果我们的生成器对象不是无限长的，我们可以将其转换为列表并访问索引 N 处的列表元素。

def g():
    yield 1
    yield 2
    yield 3


gen = g()

my_list = list(gen)

print(my_list[0])  # 👉️ 1
print(my_list[1])  # 👉️ 2
print(my_list[2])  # 👉️ 3

print(list(gen)) # 👉️ []

列表类接受一个可迭代对象并返回一个列表对象。

请注意 ，将生成器对象转换为列表会耗尽它，因此请确保将结果存储在变量中。

我们还可以使用 enumerate() 函数来访问当前迭代的索引。

def g():
    yield 1
    yield 2
    yield 3


gen = g()

nthElement = next(
    (item for index, item in enumerate(gen) if index == 1),
    None
)

print(nthElement)  # 👉️ 2