如何基于双索引数组高效切片Numpy数组（无循环实现）

刚好碰到过类似的需求，用Numpy的向量化操作就能完美解决，完全不用写Python循环，效率还比循环高很多！核心思路是先一次性生成所有需要提取的元素的全局索引，取出后再拆分成目标的子数组结构，具体步骤如下：

具体实现步骤 #

• 计算每个切片的长度：先算出index_arr1[i]:index_arr2[i]每个切片包含的元素个数，这是后续拆分结果的关键依据。
• 生成全局索引数组：用向量化操作生成所有要提取的元素在input_arr中的索引，彻底避开循环。
• 提取元素并拆分：用全局索引一次性取出所有元素，再按之前计算的切片长度拆分成子数组列表。

示例代码 #

import numpy as np

# 你的示例输入
index_arr1 = np.asarray([2,3,4])
index_arr2 = np.asarray([5,5,6])
input_arr = np.asarray([1,2,3,4,4,5,7,2])

# 1. 计算每个切片的元素个数
slice_lengths = index_arr2 - index_arr1

# 2. 生成所有需要提取的全局索引
# 重复每个起始索引，次数等于对应切片的长度
repeated_starts = np.repeat(index_arr1, slice_lengths)
# 生成每个切片内的偏移量（0到length-1）
offsets = np.concatenate([np.arange(l) for l in slice_lengths])
# 合并得到每个元素在input_arr中的索引
all_indices = repeated_starts + offsets

# 3. 一次性提取元素，再拆分成目标子数组
result = np.split(input_arr[all_indices], np.cumsum(slice_lengths)[:-1])

print(result)
# 输出：[array([3, 4, 4]), array([4, 4]), array([4, 5])]

进阶优化：彻底移除Python层面的循环 #

上面的np.concatenate([np.arange(l) for l in slice_lengths])虽然简洁，但还是有个Python层面的列表推导循环。如果你的切片数量特别多，可以用下面的方式彻底用Numpy向量化生成偏移量：

total_elements = slice_lengths.sum()
# 计算每个切片的累计起始位置（用于偏移计算）
cumulative_starts = np.cumsum([0] + slice_lengths[:-1])
# 生成全局偏移量：每个元素的偏移 = 全局位置 - 所属切片的累计起始位置
offsets = np.arange(total_elements) - np.repeat(cumulative_starts, slice_lengths)

这样所有操作都是Numpy的底层C实现，性能会更优。

为什么这个方案高效？ #

• 所有核心操作都是Numpy的内置向量化函数，比Python循环快几个数量级，尤其是处理大规模数据时差距更明显。
• 只对input_arr做一次索引提取，避免了多次切片操作的额外开销。

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