如何用NumPy向量化操作高效提取数组中指定边界内的元素？

当data规模较大时，循环切片的方法确实会因为多次内存操作和切片开销变慢，完全可以用向量化操作大幅提升效率。核心思路是先一次性提取所有非空切片的元素，再拆分并补全空数组位置，避免反复对大数组做切片操作。

具体实现步骤 #

假设我们有原数组：

import numpy as np

data = np.array([97, 98, 99, 100])
offsets = np.array([0, 1, 2, 2, 3, 3, 4])

1. 计算每个切片的长度
   先算出offsets相邻元素的差值，得到每个切片包含的元素个数：

slice_lengths = offsets[1:] - offsets[:-1]
# 结果: array([1, 1, 0, 1, 0, 1])

2. 提取所有非空切片的元素并拼接
   过滤出非空切片的起始/结束位置，一次性拼接所有需要的元素，避免多次切片大数组：

# 筛选非空切片的起止索引
non_empty_mask = slice_lengths > 0
starts = offsets[:-1][non_empty_mask]
ends = offsets[1:][non_empty_mask]

# 一次性拼接所有非空切片的元素
concatenated_elements = np.concatenate([data[s:e] for s, e in zip(starts, ends)])
# 结果: array([97, 98, 99, 100])

3. 拆分拼接后的数组并补全空数组
   根据非空切片的长度拆分拼接后的数组，再把空数组插回对应位置：

# 计算拆分位置（跳过最后一个切片，避免空拆分）
split_positions = np.cumsum(slice_lengths[non_empty_mask])[:-1]
# 拆分得到非空数组列表
non_empty_arrays = np.split(concatenated_elements, split_positions)
# 结果: [array([97]), array([98]), array([99]), array([100])]

# 补全空数组到对应位置
result = []
non_empty_idx = 0
for length in slice_lengths:
    if length > 0:
        result.append(non_empty_arrays[non_empty_idx])
        non_empty_idx += 1
    else:
        result.append(np.array([], dtype=data.dtype))

最终结果验证 #

运行后result和原循环方法的输出完全一致：

[array([97]),
 array([98]),
 array([], dtype=int64),
 array([99]),
 array([], dtype=int64),
 array([100])]

效率优势 #

这种方法把大部分操作转为向量化计算，仅最后一步用极短的循环处理空数组位置。相比原循环反复对大数组做切片，减少了大量内存分配和数组拷贝的开销，在data规模越大时，效率提升越明显。

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