如何用NumPy风格方法按可变起始索引与子序列长度提取元素

当然存在纯NumPy的实现方式，不需要依赖Python循环。核心思路是先一次性生成所有需要提取的元素索引，提取后再按指定的切片长度分割成子数组，具体实现如下：

import numpy as np

A = np.arange(1000)
x = np.array([0, 10, 20, 30], dtype=int)
dx = np.array([3, 4, 5, 6], dtype=int)

# 生成所有待提取元素的索引数组
indices = np.concatenate([np.arange(s, s+l) for s, l in zip(x, dx)])
# 批量提取元素
extracted_data = A[indices]
# 按dx指定的长度分割为子数组
result = np.split(extracted_data, np.cumsum(dx)[:-1])

# 验证输出
for arr in result:
    print(arr)

关键步骤说明： #

• 生成索引数组：通过np.concatenate结合列表推导，把每个切片的索引合并成一个一维数组，列表推导仅用于构造单段索引序列，核心合并操作由NumPy完成。
• 批量提取元素：利用NumPy的向量化索引直接从A中提取所有目标元素，效率远高于循环逐个提取。
• 分割子数组：np.cumsum(dx)[:-1]计算出分割点的位置，再用np.split把提取后的数组按指定长度分割，得到与原循环输出一致的子数组列表。

这种方式完全基于NumPy内置函数实现，符合numpythonic风格，在数据量较大时会比原Python循环有明显的性能提升。

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