如何高效在Python的NumPy数组中将指定元素替换为多个元素？

针对你需要将数组中特定元素替换为可变长度序列的需求，以下两种方法可以高效实现：

方法一：列表推导式扩展（简洁高效，适合多数场景） #

这种方法逻辑直观，代码简洁，对于常规规模的数组足够高效：

import numpy as np

arr = np.array([1,2,3,4,5,6])
# 定义替换规则：键是原数组元素，值是替换后的序列
replace_map = {4: [2,2], 6: [2,3]}

# 遍历原数组，扩展每个元素对应的序列
expanded_list = []
for num in arr:
    expanded_list.extend(replace_map.get(num, [num]))

# 转换为NumPy数组
result = np.array(expanded_list)
# 输出：array([1, 2, 3, 2, 2, 5, 2, 3])

方法二：纯NumPy向量化实现（适合超大规模数组） #

如果处理的是超大数组，想要避免Python循环带来的开销，可以用纯NumPy操作实现：

import numpy as np

arr = np.array([1,2,3,4,5,6])
replace_map = {4: np.array([2,2]), 6: np.array([2,3])}

# 计算每个元素替换后的长度
lengths = np.array([len(replace_map.get(x, [x])) for x in arr])
# 计算每个元素在结果数组中的起始索引
start_indices = np.cumsum(lengths) - lengths
# 创建结果数组
result = np.empty(np.sum(lengths), dtype=arr.dtype)

# 处理需要替换的元素
for target_val, replacement in replace_map.items():
    # 找到原数组中匹配目标值的位置
    match_mask = arr == target_val
    # 计算这些位置在结果数组中对应的填充索引
    fill_positions = start_indices[match_mask].repeat(len(replacement)) + np.tile(np.arange(len(replacement)), match_mask.sum())
    # 填充替换值
    result[fill_positions] = np.tile(replacement, match_mask.sum())

# 处理无需替换的元素
non_match_mask = ~np.isin(arr, list(replace_map.keys()))
result[start_indices[non_match_mask]] = arr[non_match_mask]

# 输出：array([1, 2, 3, 2, 2, 5, 2, 3])

关于np.select的局限性 #

np.select的设计目标是根据条件选择等长度的输出值，无法处理将单个元素替换为多元素序列的场景，因此不适合你的需求。

内容的提问来源于stack exchange，提问作者Matias Bay Rodriguez