How to create a 2D array of ranges using numpy

Question

I have an array of start and stop indices, like this:

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

and i would like to construct a 2D numpy array where each row is a range from the corresponding pair of start and stop indices, as follows:

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

Currently, i am creating an empty array and filling it in a for loop:

    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])

Is there a more compact and (more importantly) faster way of doing this? possibly something that doesn't involve using a loop?

Just do `np.asarray(a)[:, :1] + np.arange(a[0][1] - a[0][0])` — jdehesa, May 16 '19 at 16:32
Would the difference between start and stop across all pairs be the same? — Divakar, May 16 '19 at 16:34

Andy Hayden · Accepted Answer · 2019-05-16T16:40:59.993

4

One way is to use broadcast to add the left hand edges to the base arange:

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

Note: this requires all ranges to be the same length.

edited May 16 '19 at 16:40

answered May 16 '19 at 16:30

Andy Hayden

Awesome! thanks. @jdehesa's solution also works but this seems to run a bit faster – Jonathan May 17 '19 at 08:48

yatu · Answer 2 · 2019-05-17T07:42:58.380

2

If the ranges were to result in different lengths, for a vectorized approach you could use n_ranges from the linked solution:

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])]

Which would also work with the following array:

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])]

edited May 17 '19 at 07:42

answered May 16 '19 at 16:28

yatu

Thanks for the complete solution, but i forgot to make it clear in my question that the ranges would all be the same length. – Jonathan May 17 '19 at 09:14

2 Answers2