CODEWITHSUNDEEP
pythonoptimizationnumpy
jterrace
jterrace

Reputation: 67063

Array of ranges with numpy

I have the following array:

>>> x = numpy.array([2,4,2,3,1])
>>> x
array([2, 4, 2, 3, 1])

I would like an array of ranges of these values. I can create it like this:

>>> numpy.hstack( (numpy.arange(v) for v in x) )
array([0, 1, 0, 1, 2, 3, 0, 1, 0, 1, 2, 0])

Given x, is there a faster way to generate this with numpy without having to use a for loop?

Upvotes: 2

Views: 2539

Answers (2)

Yann
Yann

Reputation: 671

Is this actually faster ? I have a similar need, and

concatenate([range(l, r) for l, r in array((left, right)).T])

is twice as fast as

range(end[-1]) + repeat(left + end, right-left)

(where end = cumsum(right - left) just like yours).

(in my very short experience, repeat is very slow -- at least in python 3.6)

Upvotes: 0

jterrace
jterrace

Reputation: 67063

I figured it out:

>>> x
array([2, 4, 2, 3, 1])
>>> ends = numpy.cumsum(x)
>>> ranges = numpy.arange(ends[-1])
>>> ranges = ranges - numpy.repeat(ends-x, x)
>>> ranges
array([0, 1, 0, 1, 2, 3, 0, 1, 0, 1, 2, 0])
>>>

Upvotes: 1

Related Questions