Range of all elements in numpy array

Question

Suppose I have a 1d array a where from each element I would like to have a range of which the size is stored in ranges:

a = np.array([10,9,12])
ranges = np.array([2,4,3])

The desired output would be:

np.array([10,11,9,10,11,12,12,13,14])

I could of course use a for loop, but I prefer a fully vectorized approach. np.repeat allows one to repeat the elements in a a number of times by setting repeats=, but I am not aware of a similar numpy function particularly dealing with the problem above.

This seems to be enough of a niche occurrence that I can imagine that no specific functionality exists for this use case in NumPy. hstack with a list comprehension involving arange seems to be the most straightforward solution. — 9769953
– 9769953, Commented Sep 6, 2021 at 8:52
Are you using this often (e.g., in a tight loop), that you prefer a vectorized approach? For a one-off (or few times usage), any solution that works will be fine and not hamper performance. — 9769953
– 9769953, Commented Sep 6, 2021 at 8:55
Hmm, in principle all the information needed is np.repeat(a, ranges), as it will produce array([10, 10, 9, 9, 9, 9, 12, 12, 12]) . I am just a bit stuck on how to convert the consecutive repeats to ranges — pr94
– pr94, Commented Sep 6, 2021 at 8:57
It is certainly something I would need recurrently, but perhaps you're right a loop is still good enough. I am just always trying to avoid as much as possiblethem for code I'll over and over. Hence the question asked here — pr94
– pr94, Commented Sep 6, 2021 at 9:02

9769953 · Accepted Answer · 2021-09-06 09:08:35Z

2

>>> np.hstack([np.arange(start, start+size) for start, size in zip(a, ranges)])
array([10, 11,  9, 10, 11, 12, 12, 13, 14])

edited Sep 6, 2021 at 9:08

answered Sep 6, 2021 at 9:03

9769953

12.5k3 gold badges31 silver badges46 bronze badges

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

Comments

U13-Forward · Accepted Answer · 2021-09-06 09:00:32Z

1

With pandas it could be easier:

>>> import pandas as pd
>>> x = pd.Series(np.repeat(a, ranges))
>>> x + x.groupby(x).cumcount()
0    10
1    11
2     9
3    10
4    11
5    12
6    12
7    13
8    14
dtype: int64
>>>

If you want a numpy array:

>>> x.add(x.groupby(x).cumcount()).to_numpy()
array([10, 11,  9, 10, 11, 12, 12, 13, 14], dtype=int64)
>>>

answered Sep 6, 2021 at 9:00

U13-Forward

71.8k15 gold badges100 silver badges125 bronze badges

11 Comments

pr94 Over a year ago

Certainly not a bad solution. I wonder how it compares to list comprehension in terms of timing

U13-Forward Over a year ago

@pr94 This would be much quicker than list comprehension.

U13-Forward Over a year ago

@pr94 Numpy is developed with C, list comprehensions are slow, so Using this might be hundreds of times faster

pr94 Over a year ago

Yes I know. I thought perhaps the use of pandas would slow down a bit

U13-Forward Over a year ago

@pr94 Pandas backend is in Numpy, so it's the same speed as numpy

|

9769953 · Accepted Answer · 2021-09-06 09:45:29Z

1

Someone asked about timing, so I compared the times of the three solutions (so far) in a very simple manner, using the %timeit magic function in Jupyter notebook cells.

I set it up as follows:

N = 1
a = np.array([10,9,12])
a = np.tile(a, N)
ranges = np.array([2,4,3])
ranges = np.tile(ranges, N)
a.shape, ranges.shape

So I could easily scale (albeit things not random, but repeated).

Then I ran:

%timeit np.hstack([np.arange(start, start+size) for start, size in zip(a, ranges)])

,

%timeit x = pd.Series(np.repeat(a, ranges)); x.add(x.groupby(x).cumcount()).to_numpy()

and

%timeit np.array([i for j in range(len(a)) for i in range(a[j],a[j]+ranges[j])])

Results are as follows: N = 1:

9.81 µs ± 481 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
568 µs ± 20.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
3.53 µs ± 81.4 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

N = 10:

63.4 µs ± 976 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
575 µs ± 15.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
25.1 µs ± 698 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

N = 100:

612 µs ± 12.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
608 µs ± 25.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
237 µs ± 9.62 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

N = 1000:

6.09 ms ± 52 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
852 µs ± 2.66 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
2.44 ms ± 43.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

So the Pandas solution wins when things get to arrays of 1000 elements or more, but the Python double list comprehension does an excellent job until that point. np.hstack probably loses out because of extra memory allocation and copying, but that's a guess. Note also that the Pandas solution is nearly the same time for each array size.

Caveats still exists because there are repeated numbers, and all values are relatively small integers. This really shouldn't matter, but I'm not (yet) betting on it. (For example, Pandas groupby functionality may be fast because of the repeated numbers.)

Bonus: the OP has statement in a comment that "The real life arrays are around 1000 elements, yet with ranges ranging from 100 to 1000. So becomes quite big – pr94". So I adjusted my timing test to the following:

import numpy as np
import pandas as pd

N = 1000
a = np.random.randint(100, 1000, N)
# This is how I understand "ranges ranging from 100 to 1000"
ranges = np.random.randint(100, 1000, N)

%timeit np.hstack([np.arange(start, start+size) for start, size in zip(a, ranges)])
%timeit  x = pd.Series(np.repeat(a, ranges)); x.add(x.groupby(x).cumcount()).to_numpy()
%timeit np.array([i for j in range(len(a)) for i in range(a[j],a[j]+ranges[j])])

Which comes out as :

hstack: 2.78 ms ± 38.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
pandas: 18.4 ms ± 663 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
double list comprehension: 64.1 ms ± 427 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Which shows that those caveats I mentioned, in some form at least, do seem to exist. But people should double check whether this testing code is actually the most relevant and appropriate, and whether it is correct.

edited Sep 6, 2021 at 9:45

answered Sep 6, 2021 at 9:19

9769953

12.5k3 gold badges31 silver badges46 bronze badges

16 Comments

9769953 Over a year ago

@pr94 That is up to you, and really depends what you prefer: introducing another dependency (Pandas, although you may already be using it), readability (I have to read the Pandas solution several times to understand it; np.hstack with np.arange is very straightforward to read) or just plain speed. Also, speed hardly means anything if this part of your code is a few percent of less in time compared to the overall code.

pr94 Over a year ago

You're absolutely right that a list comprehension is more readable

9769953 Over a year ago

Updated with %timeit, so better timings (I hope).

Mahdi F. Over a year ago

@pr94 this is not correct you say I select pandas better is you say base on your work decision to select which solution if your array is short you can see different runtime that list is best and when your array is big you need pandas and is better.

9769953 Over a year ago

@user1740577 I have a hard time understanding your comment, but it seems to be about the fact that the best solution depends on the size. There is, however, a comment (sadly not in the question itself) from the OP staing their arrays are 1000s of elements, with ranges ranging from 100 to 1000 ("The real life arrays are around 1000 elements, yet with ranges ranging from 100 to 1000. So becomes quite big – pr94"). This is why I went with increasing numbers up to 1000 elements.

|

norok2 · Accepted Answer · 2021-09-07 08:51:06Z

This problem is probably going to be solved much faster with a Numba-compiled function:

@nb.jit 
def expand_range(values, counts): 
    n = len(values) 
    m = np.sum(counts) 
    r = np.zeros((m,), dtype=values.dtype) 
    k = 0 
    for i in range(n): 
        x = values[i] 
        for j in range(counts[i]): 
            r[k] = x + j 
            k += 1 
    return r

On the very small inputs:

%timeit expand_range(a, ranges)                                                                                                                                                
# 1.16 µs ± 126 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

%timeit x = pd.Series(np.repeat(a, ranges)); x.add(x.groupby(x).cumcount()).to_numpy()                                                                                         
# 617 µs ± 4.32 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit np.hstack([np.arange(start, start+size) for start, size in zip(a, ranges)])                                                                                            
# 25 µs ± 2.2 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit np.array([i for j in range(len(a)) for i in range(a[j],a[j]+ranges[j])])                                                                                               
# 13.5 µs ± 929 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

and on somewhat larger inputs:

b = np.random.randint(0, 1000, 1000)
b_ranges = np.random.randint(1, 10, 1000)

%timeit expand_range(b, b_ranges)                                                                                                                                              
# 5.07 µs ± 98.1 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit x = pd.Series(np.repeat(a, ranges)); x.add(x.groupby(x).cumcount()).to_numpy()                                                                                         
# 617 µs ± 4.32 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit np.hstack([np.arange(start, start+size) for start, size in zip(a, ranges)])                                                                                            
# 25 µs ± 2.2 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit np.array([i for j in range(len(a)) for i in range(a[j],a[j]+ranges[j])])                                                                                               
# 13.5 µs ± 929 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

these show that with Numba-based approach winning the speed gain is at least 100x over any of the other approaches proposed so far.

With the numbers closer to what as been indicated in one of the comments by the OP:

b = np.random.randint(10, 1000, 1000)
b_ranges = np.random.randint(100, 1000, 1000)

%timeit expand_range(b, b_ranges)                                                                                                                                              
# 1.5 ms ± 67.9 µs per loop (mean ± std. dev. of 7 runs, 1000

%timeit x = pd.Series(np.repeat(b, b_ranges)); x.add(x.groupby(x).cumcount()).to_numpy()                                                                                       
# 91.8 ms ± 6.53 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit np.hstack([np.arange(start, start+size) for start, size in zip(b, b_ranges)])                                                                                          
# 10.7 ms ± 402 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit np.array([i for j in range(len(b)) for i in range(b[j],b[j]+b_ranges[j])])                                                                                             
# 144 ms ± 4.54 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

which is still at least a respectable 7x over the others.

Collectives™ on Stack Overflow

Range of all elements in numpy array

4 Answers 4

Comments

11 Comments

16 Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

4 Answers 4

Comments

11 Comments

16 Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related