如何将如下结构优化为向量化操作以提升运行效率？

原代码依赖Python循环+tile+concatenate生成c_indices，当_c规模较大时（比如10000个元素），循环和多次内存分配会导致速度严重下降。可以用完全向量化的numpy操作彻底解决这个问题，以下是实现方案：

向量化实现代码 #

import numpy as np

_c = np.random.randint(0, 1000, (10000))

# 计算每个分段的起始位置
c_starts = np.concatenate([[0], np.cumsum(_c[:-1])])
# 计算每个分段的总元素数（分段长度 × 重复次数）
total_per_block = _c * _c
# 累计每个分段的元素数，用于定位全局位置所属的分段
cum_total = np.cumsum(total_per_block)
# 生成全局位置数组
pos = np.arange(cum_total[-1])
# 找到每个全局位置对应的分段索引
block_indices = np.searchsorted(cum_total, pos, side='right')
# 计算当前位置在所属分段内的偏移量
block_offsets = pos - np.concatenate([[0], cum_total[:-1]])[block_indices]
# 计算分段内的循环偏移值（实现tile分段的效果）
k = block_offsets % _c[block_indices]
# 拼接起始位置和偏移值得到最终结果
c_indices = c_starts[block_indices] + k

原理说明 #

这个方案通过以下步骤实现完全向量化：

1. 先计算每个分段的起始位置和总元素数；
2. 用全局位置数组pos覆盖所有最终元素的位置；
3. 通过searchsorted快速定位每个位置所属的分段；
4. 利用取模运算%模拟tile分段的效果（每个分段的元素序列重复指定次数）；
5. 最后将分段起始位置与偏移值相加，得到最终的c_indices。

正确性验证 #

以_c = [3,4,2]为例，上述代码生成的c_indices与预期结果完全一致：

• 分段[0,1,2]重复3次；
• 分段[3,4,5,6]重复4次；
• 分段[7,8]重复2次。

性能优势 #

原代码的Python循环会带来大量额外开销，而向量化方案所有操作都是numpy的C级内置实现，对于10000元素的_c，速度能提升一个数量级以上，且内存分配更高效。

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