如何基于Numpy高效实现数组索引映射与对应值扩展？

问题描述 #

需要将一维Numpy数组overlay的元素映射到更大的一维数组arr上：overlay的每个元素对应arr中的一组索引（存储在subixs中），修改overlay的元素时，要同步更新arr中对应索引的所有值。

示例场景：

import numpy as np

overlay = np.array([0, 1, 1, 4, 3])
subixs = [[0, 1, 2, 3], [4, 5], [6], [7], [8]]  # overlay[i]对应arr的subixs[i]索引组
arr = np.array([0, 0, 0, 0, 1, 1, 1, 4, 3])

• 修改overlay[1] = 7时，arr需变为[0, 0, 0, 0, 7, 7, 1, 4, 3]
• 当用布尔索引批量修改时：

  msk = (overlay == 1) | (overlay == 3)
  overlay[msk] = [44, 48, 47]
  

  需要将msk转换为arr的索引[4,5,6,8]，同时把值[44,48,47]扩展为[44,44,48,47]后更新arr。

现有实现依赖列表拼接（如sum(subixs[msk], [])），在arr元素达10万级时效率低下，需用Numpy向量化操作优化。

高效解决方案 #

核心思路是提前预处理映射关系，将subixs转换为Numpy原生数组格式，避免每次操作都做列表级别的拼接：

步骤1：预处理映射结构 #

一次性将subixs扁平化，并生成用于关联overlay和arr的辅助数组：

# 将subixs转换为扁平化的arr索引数组
flat_subixs = np.concatenate(subixs)
# 记录每个overlay元素对应的arr索引数量
repeat_counts = np.array([len(ix) for ix in subixs])
# 生成重复的overlay索引数组：每个overlay[i]对应repeat_counts[i]个自身索引
overlay_indices_repeated = np.repeat(np.arange(len(overlay)), repeat_counts)

步骤2：布尔索引的批量更新 #

当需要通过msk批量修改overlay并同步arr时：

msk = (overlay == 1) | (overlay == 3)
new_values = np.array([44, 48, 47])
overlay[msk] = new_values

# 筛选出需要更新的arr索引
target_overlay_indices = np.flatnonzero(msk)
update_arr_indices = flat_subixs[np.isin(overlay_indices_repeated, target_overlay_indices)]
# 按对应索引组的长度扩展新值
expanded_values = np.repeat(new_values, repeat_counts[msk])

# 批量更新arr
arr[update_arr_indices] = expanded_values

步骤3：单个元素的快速更新 #

如果仅修改单个overlay元素（如overlay[1] = 7），直接利用预处理的数组快速定位：

idx = 1
new_val = 7
overlay[idx] = new_val
# 计算该元素对应的arr索引切片范围
start_idx = repeat_counts[:idx].sum()
end_idx = start_idx + repeat_counts[idx]
# 更新arr对应位置
arr[flat_subixs[start_idx:end_idx]] = new_val

方案优势 #

• 预处理仅需一次，后续所有更新操作均为纯Numpy向量化计算，彻底避免列表拼接的性能开销
• 针对10万级规模的arr，效率远高于基于列表操作的实现
• 同时支持布尔索引批量更新和单个索引快速更新两种场景

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