Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

How to create a 2D array of ranges using numpy

Tags:

python

numpy

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?

like image 453
Jonathan Avatar asked May 16 '19 16:05

Jonathan


1 Answers

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.

like image 148
Andy Hayden Avatar answered Oct 19 '22 04:10

Andy Hayden