如何用numpy.arange以两个数组为起止参数生成指定二维数组？

问题描述 #

需要构建一个二维数组，满足以下规则：

• 第一维度的长度与输入的arr_1、arr_2两个数组一致；
• 每个子数组由numpy.arange生成：以arr_1的对应元素为start参数，以arr_2对应元素加1作为stop参数（因为arange是左闭右开区间，要包含arr_2对应元素，需将stop设为arr_2[i]+1）。

输入示例：

import numpy as np
arr_1 = np.array([0, 0, 0, 1, 2, 3, 4])
arr_2 = np.array([0, 1, 2, 3, 4, 5, 6])

期望输出：

res_arr = np.array([[0], [0, 1], [0, 1, 2], [1, 2, 3], [2, 3, 4], [3, 4, 5], [4, 5, 6]])

解决方案 #

可以通过NumPy实现，核心是利用NumPy的object类型数组容纳长度可变的子数组，以下是两种实用实现方式：

方法1：列表推导式（直观易读） #

直接遍历两个数组的对应元素生成子数组，再转为object类型的NumPy数组：

import numpy as np

arr_1 = np.array([0, 0, 0, 1, 2, 3, 4])
arr_2 = np.array([0, 1, 2, 3, 4, 5, 6])

res_arr = np.array([np.arange(start, stop + 1) for start, stop in zip(arr_1, arr_2)], dtype=object)

方法2：向量化分割（进阶） #

通过生成所有元素的一维数组，再按子数组长度分割，避免显式循环：

import numpy as np

arr_1 = np.array([0, 0, 0, 1, 2, 3, 4])
arr_2 = np.array([0, 1, 2, 3, 4, 5, 6])

# 计算每个子数组的长度
lengths = arr_2 - arr_1 + 1
# 生成所有元素的一维数组
all_elements = np.repeat(arr_1, lengths) + np.arange(lengths.sum())
# 分割为子数组并转为object数组
res_arr = np.array(np.split(all_elements, np.cumsum(lengths)[:-1]), dtype=object)

注意：由于子数组长度不一致，生成的是dtype=object的数组，这类数组无法像常规同形二维数组那样进行全量向量化运算，若后续有大量数值计算需求，建议根据场景调整数据结构。

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