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]) — javidcf
– javidcf, Commented May 16, 2019 at 16:32
Would the difference between start and stop across all pairs be the same? — Divakar
– Divakar, Commented May 16, 2019 at 16:34

Andy Hayden · Accepted Answer · 2019-05-16 16:40:59Z

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.

answered May 16, 2019 at 16:30

Andy Hayden

378k110 gold badges640 silver badges546 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

Awesome! thanks. @jdehesa's solution also works but this seems to run a bit faster

yatu · Accepted Answer · 2019-05-17 07:42:58Z

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

answered May 16, 2019 at 16:28

yatu

88.6k12 gold badges93 silver badges148 bronze badges

Thanks for the complete solution, but i forgot to make it clear in my question that the ranges would all be the same length.

@yatu I don't see the return_flat-variable in the linked code?