Python 多列表按轮换规则选取元素的更简洁实现方案咨询

首先明确核心规则：第k次选取（k从0开始计数）时，取第k % 列表总数个列表的第k位元素，持续选取直到任意一个列表索引越界。

原有代码可优化点 #

• 硬编码4个输入文件参数，无法适配更多/更少列表的场景
• 多分支判断逻辑冗余，维护成本高
• 边界判断依赖单个列表长度，若其他列表更短会触发索引报错
• 代码与调试打印语句耦合，生产使用需要额外清理

优化实现 #

通用可扩展版本 #

支持任意数量的输入列表，拆分生成器和业务逻辑，可按需取数：

from typing import List, Set, Generator

# 生成器：可无限按规则生成元素（只要对应列表有足够长度）
def rotated_picker(*files: List) -> Generator:
    k = 0
    while True:
        target_file_idx = k % len(files)
        yield files[target_file_idx][k]
        k += 1

# 封装方法：取到所有列表能支持的最大数量元素，返回集合
def get_rotated(*files: List) -> Set:
    # 取所有列表的最短长度作为最大选取次数
    max_pick_count = min(len(f) for f in files)
    picker = rotated_picker(*files)
    return {next(picker) for _ in range(max_pick_count)}

测试验证 #

使用提供的测试数据运行：

import numpy as np
file1 = np.arange(0, 4).tolist()
file2 = np.arange(3, 7).tolist()
file3 = np.arange(7, 11).tolist()
file4 = np.arange(11, 15).tolist()

print(get_rotated(file1, file2, file3, file4))
# 输出：{0, 4, 9, 14}，与预期结果完全一致

极简版本 #

如果确定是4个列表，且固定取4个元素，可直接用集合推导式一行实现：

files = [file1, file2, file3, file4]
final = {files[k%4][k] for k in range(4)}

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