如何不使用for循环按边界数组分割NumPy一维数组？

给定一维NumPy数组：

import numpy as np

a = np.array([1,2,3,5,6,8,9,11,12,13,14,15])
b = np.array([0,3,6,9,12,15,18])

需要将a按b定义的区间分割（第i个子数组包含b[i]与b[i+1]之间的元素），最终得到：

result = [[1,2],
          [3,5],
          [6,8],
          [9,11],
          [12,13,14],
          [15]]

高效实现方案 #

利用NumPy的向量化操作替代循环，核心是用np.digitize快速定位每个元素所属的区间，再通过np.split完成分割：

1. 标记元素所属区间
   用np.digitize获取数组a中每个元素对应的区间索引，通过right=False确保区间为左闭右开（匹配题目要求的[b[i], b[i+1})规则）：

indices = np.digitize(a, b, right=False)

执行后indices为array([1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 5, 6])，每个数字代表对应元素所属的区间序号。

2. 定位分割点并分割数组
   通过np.diff找到索引变化的位置，以此确定分割点，再用np.split完成分割：

# 找到索引发生变化的位置，计算分割点
split_points = np.where(np.diff(indices))[0] + 1
# 分割数组
result = np.split(a, split_points)

此时result即为所需的子数组列表，若需要转为Python列表格式，直接用list(result)即可。

另一种等价实现 #

也可以通过np.unique统计每个区间的元素数量，再用累积和计算分割点：

_, counts = np.unique(indices, return_counts=True)
result = np.split(a, np.cumsum(counts)[:-1])

效率说明 #

以上两种方案均为NumPy内置的向量化操作，完全避免了Python层面的循环，在处理大规模数组时，速度比for循环快几个数量级。

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