CODEWITHSUNDEEP
pythonnumpy
Jonathan
Jonathan

Reputation: 53

How to create a 2D array of ranges using 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?

Upvotes: 5

Views: 2582

Answers (2)

yatu
yatu

Reputation: 88226

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

Upvotes: 2

Andy Hayden
Andy Hayden

Reputation: 375375

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.

Upvotes: 4

Related Questions