调整Python/NumPy代码 实现任意输入下均输出列表的列表避免迭代报错

核心逻辑是统一两种分支的输出格式，确保所有情况下Z的每个元素都是单层列表，同时避免numpy object数组的自动嵌套问题，具体实现如下：

实现代码 #

import numpy as np

# 输入可根据需求动态调整
Arr1 = np.array([9,7,3,1])
Arr2 = [[14,6],[13,2]]
Arr3 = np.array([0,2])

# 构建索引替换映射，提升查询效率
replace_mapping = {index: Arr2[pos] for pos, index in enumerate(Arr3)}

# 生成目标列表，所有元素均为可迭代的列表类型
Z = [replace_mapping[i] if i in replace_mapping else [Arr1[i]] for i in range(len(Arr1))]

兼容效果验证 #

• 当Arr3非空时，输出Z = [[14, 6], [7], [13, 2], [1]]，每个元素都是单层列表，可直接遍历
• 当Arr3为空时，输出Z = [[9], [7], [3], [1]]，所有元素为单元素列表，不会出现不可迭代的数值类型

原问题原因说明 #

1. 原有第一种写法将生成的列表传入np.array指定dtype=object时，numpy会对嵌套结构做自动封装，导致本来就是列表的Arr2元素被额外套了一层列表，迭代时出现类型不匹配
2. 第二种写法去掉[Arr1[i]]的外层括号后，未匹配的分支直接返回数值类型，迭代单个数值时触发不可迭代报错

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