Question

假设我想创建一个从列表中返回固定数量用户的生成器，如果我想要的用户数大于列表的长度，它只会重复最后一个用户。我写了一个生成器来执行以下操作：

def gen_users(users, total):
    num_users = len(users)
    cur_user = 0
    while cur_user < total:
        if cur_user >= num_users:
            yield users[-1]
        else:
            yield users[cur_user]
        cur_user += 1

因此，如果我有一个用户列表users = ['one', 'two', 'three']并且我致电gen_users(users, 5)，我会回来one, two, three, three, three。

有更多的pythonic方法吗？使用生成器是解决此问题的有效方法吗？

Answer 1

使用itertools.islice：

>>> import itertools
>>> 
>>> def gen_users(users, total):
...     # assert users and total > 0
...     for i, user in enumerate(itertools.islice(users, total)):
...         yield user
...     for _ in range(total - i - 1): # i -> last index
...         yield user
... 
>>> list(gen_users([1,2,3], 5))
[1, 2, 3, 3, 3]
>>> list(gen_users([1,2,3], 2))
[1, 2]

或者您可以使用users[:total]，但它会创建临时序列。

<强>更新

Hans Zauber代码的略微修改版本：

>>> from itertools import islice, chain, repeat
>>> 
>>> def gen_users(users, total):
...     return islice(chain(users, repeat(users[-1])), total)
... 
>>> list(gen_users([1,2,3], 5))
[1, 2, 3, 3, 3]
>>> list(gen_users([1,2,3], 3))
[1, 2, 3]
>>> list(gen_users([1,2,3], 2))
[1, 2]

Answer 2

试试这个单行：

import itertools
gen_users = lambda users, total: itertools.chain(itertools.islice(users, total), itertools.repeat(users[-1], total-len(users)))

有没有更好的方法来定义这个发电机？

2 个答案: