如何使用numpy創建范圍的2D數組

Question

我有一系列啟動和停止索引，如下所示：

    [[0, 3], [4, 7], [15, 18]]

我想構建一個2D numpy數組，其中每一行都是相應的一對起始和終止索引的范圍，如下所示：

    [[0, 1, 2],
    [4, 5, 6],
    [15, 16, 18]]

目前，我正在創建一個空數組並在for循環中填充它：

    ranges = numpy.empty((3, 3))
    a = [[0, 3], [4, 7], [15, 18]]

    for i, r in enumerate(a):
        ranges[i] = numpy.arange(r[0], r[1])

有更緊湊和（更重要的）更快的方式嗎？ 可能是一些不涉及使用循環的東西？

Answer 1

一種方法是使用廣播將左手邊緣添加到基本范圍：

In [11]: np.arange(3) + np.array([0, 4, 15])[:, None]
Out[11]:
array([[ 0,  1,  2],
       [ 4,  5,  6],
       [15, 16, 17]])

注意：這要求所有范圍都是相同的長度。

Answer 2

如果范圍導致不同的長度，對於矢量化方法，您可以使用鏈接解決方案中的n_ranges ：

a = np.array([[0, 3], [4, 7], [15, 18]])

n_ranges(a[:,0], a[:,1], return_flat=False)
# [array([0, 1, 2]), array([4, 5, 6]), array([15, 16, 17])]

---

哪個也適用於以下數組：

a = np.array([[0, 3], [4, 9], [15, 18]])
n_ranges(*a.T, return_flat=False)
# [array([0, 1, 2]), array([4, 5, 6, 7, 8]), array([15, 16, 17])]