如何为数组生成多索引掩码？单索引代码泛化需求

修改思路 #

原代码仅支持单个索引匹配，要扩展为多索引，只需将判断条件从idx == rhosub改为idx 是否包含在目标索引集合中，同时兼容单索引输入场景。

优化后的代码（兼容单/多索引） #

显式循环版本（贴近原代码逻辑）

rhosub = [0, 1, 3]  # 可传入单个数字（如0）或多索引列表/元组
# 转换为集合提升查找效率，尤其当索引数量较多时
rhosub_set = set(rhosub) if isinstance(rhosub, (list, tuple)) else {rhosub}
mask = []
for idx, n in enumerate(rho_ss.dims[0]):
    mask.append(idx in rhosub_set)
mask

简洁列表推导式版本

rhosub = [0, 1, 3]
rhosub_set = set(rhosub) if isinstance(rhosub, (list, tuple)) else {rhosub}
mask = [idx in rhosub_set for idx, _ in enumerate(rho_ss.dims[0])]

说明 #

• 当传入rhosub=[0,1,3]时，生成的mask会在索引0、1、3位置为True，其余为False，完全符合需求；
• 传入单个索引（如rhosub=0）时，生成结果与原代码一致，兼容原有使用场景；
• 将目标索引转为集合是为了优化in操作的性能，数据量越大，性能提升越明显。

内容的提问来源于stack exchange，提问作者J.Agusti