Build numpy array with multiple custom index ranges without explicit loop

Question

In Numpy, is there a pythonic way to create array3 with custom ranges from array1 and array2 without a loop? The straightforward solution of iterating over the ranges works

Answer 1

Assuming the ranges do not overlap, you could build a mask which is nonzero where the index is between the ranges specified by array1 and array2 and then use np.flatnonzero to obtain an array of indices -- the desired array3:

import numpy as np

array1 = np.array([10, 65, 200]) 
array2 = np.array([14, 70, 204])

first, last = array1.min(), array2.max()
array3 = np.zeros(last-first+1, dtype='i1')
array3[array1-first] = 1
array3[array2-first] = -1
array3 = np.flatnonzero(array3.cumsum())+first
print(array3)

yields

[ 10  11  12  13  65  66  67  68  69 200 201 202 203]

---

For large len(array1), using_flatnonzero can be significantly faster than using_loop:

def using_flatnonzero(array1, array2):
    first, last = array1.min(), array2.max()
    array3 = np.zeros(last-first+1, dtype='i1')
    array3[array1-first] = 1
    array3[array2-first] = -1
    return np.flatnonzero(array3.cumsum())+first

def using_loop(array1, array2):
    return np.concatenate([np.arange(array1[i], array2[i]) for i in
                           np.arange(0,len(array1))])


array1, array2 = (np.random.choice(range(1, 11), size=10**4, replace=True)
                  .cumsum().reshape(2, -1, order='F'))


assert np.allclose(using_flatnonzero(array1, array2), using_loop(array1, array2))

---

In [260]: %timeit using_loop(array1, array2)
100 loops, best of 3: 9.36 ms per loop

In [261]: %timeit using_flatnonzero(array1, array2)
1000 loops, best of 3: 564 µs per loop

---

If the ranges overlap, then using_loop will return an array3 which contains duplicates. using_flatnonzero returns an array with no duplicates.

---

Explanation: Let's look at a small example with

array1 = np.array([10, 65, 200]) 
array2 = np.array([14, 70, 204])

The objective is to build an array which looks like goal, below. The 1's are located at index values [ 10, 11, 12, 13, 65, 66, 67, 68, 69, 200, 201, 202, 203] (i.e. array3):

In [306]: goal
Out[306]: 
array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1,
       1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1], dtype=int8)

Once we have the goal array, array3 can be obtained with a call to np.flatnonzero:

In [307]: np.flatnonzero(goal)
Out[307]: array([ 10,  11,  12,  13,  65,  66,  67,  68,  69, 200, 201, 202, 203])

goal has the same length as array2.max():

In [308]: array2.max()
Out[308]: 204

In [309]: goal.shape
Out[309]: (204,)

So we can begin by allocating

goal = np.zeros(array2.max()+1, dtype='i1')

and then filling in 1's at the index locations given by array1 and -1's at the indices given by array2:

In [311]: goal[array1] = 1
In [312]: goal[array2] = -1
In [313]: goal
Out[313]: 
array([ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  1,  0,  0,  0, -1,  0,  0,
        0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
        0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
        0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  1,  0,  0,
        0,  0, -1,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
        0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
        0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
        0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
        0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
        0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
        0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
        0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  1,  0,  0,  0,
       -1], dtype=int8)

Now applying cumsum (the cumulative sum) produces the desired goal array:

In [314]: goal = goal.cumsum(); goal
Out[314]: 
array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1,
       1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 0])

In [315]: np.flatnonzero(goal)
Out[315]: array([ 10,  11,  12,  13,  65,  66,  67,  68,  69, 200, 201, 202, 203])

That's the main idea behind using_flatnonzero. The subtraction of first was simply to save a bit of memory.

Answer 2

Prospective Approach

I will go backwards on how to approach this problem.

Take the sample listed in the question. We have -

array1 = np.array([10, 65, 200])
array2 = np.array([14, 70, 204])

Now, look at the desired result -

result: [10,11,12,13,65,66,67,68,69,200,201,202,203]

Let's calculate the group lengths, as we would be needing those to explain the solution approach next.

In [58]: lens = array2 - array1

In [59]: lens
Out[59]: array([4, 5, 4])

The idea is to use 1's initialized array, which when cumumlative summed across the entire length would give us the desired result. This cumumlative summation would be the last step to our solution. Why 1's initialized? Well, because we have an array that increasing in steps of 1's except at specific places where we have shifts corresponding to new groups coming in.

Now, since cumsum would be the last step, so the step before it should give us something like -

array([ 10,   1,   1,   1,  52,   1,   1,   1,   1, 131,   1,   1,   1])

As discussed before, it's 1's filled with [10,52,131] at specific places. That 10 seems to be coming in from the first element in array1, but what about the rest? The second one 52 came in as 65-13 (looking at the result) and in it 13 came in the group that started with 10 and ran because of the length of the first group 4. So, if we do 65 - 10 - 4, we will get 51 and then add 1 to accomodate for boundary stop, we would have 52, which is the desired shifting value. Similarly, we would get 131.

Thus, those shifting-values could be computed, like so -

In [62]: np.diff(array1) - lens[:-1]+1
Out[62]: array([ 52, 131])

Next up, to get those shifting-places where such shifts occur, we can simply do cumulative summation on the group lengths -

In [65]: lens[:-1].cumsum()
Out[65]: array([4, 9])

For completeness, we need to pre-append 0 with the array of shifting-places and array1[0] for shifting-values.

So, we are set to present our approach in a step-by-step format!

---

Putting back the pieces

1] Get lengths of each group :

lens = array2 - array1

2] Get indices at which shifts occur and values to be put in 1's initialized array :

shift_idx = np.hstack((0,lens[:-1].cumsum()))
shift_vals = np.hstack((array1[0],np.diff(array1) - lens[:-1]+1))

3] Setup 1's initialized ID array for inserting those values at those indices listed in the step before :

id_arr = np.ones(lens.sum(),dtype=array1.dtype)
id_arr[shift_idx] = shift_vals

4] Finally do cumulative summation on the ID array :

output = id_arr.cumsum() 

Listed in a function format, we would have -

def using_ones_cumsum(array1, array2):
    lens = array2 - array1
    shift_idx = np.hstack((0,lens[:-1].cumsum()))
    shift_vals = np.hstack((array1[0],np.diff(array1) - lens[:-1]+1))
    id_arr = np.ones(lens.sum(),dtype=array1.dtype)
    id_arr[shift_idx] = shift_vals
    return id_arr.cumsum()
  

And it works on overlapping ranges too!

In [67]: array1 = np.array([10, 11, 200]) 
    ...: array2 = np.array([14, 18, 204])
    ...: 

In [68]: using_ones_cumsum(array1, array2)
Out[68]: 
array([ 10,  11,  12,  13,  11,  12,  13,  14,  15,  16,  17, 200, 201,
       202, 203])

---

Runtime test

Let's time the proposed approach against the other vectorized approach in @unutbu's flatnonzero based solution, which already proved to be much better than the loopy approach -

In [38]: array1, array2 = (np.random.choice(range(1, 11), size=10**4, replace=True)
    ...:                   .cumsum().reshape(2, -1, order='F'))

In [39]: %timeit using_flatnonzero(array1, array2)
1000 loops, best of 3: 889 µs per loop

In [40]: %timeit using_ones_cumsum(array1, array2)
1000 loops, best of 3: 235 µs per loop

---

Improvement!

Now, codewise NumPy doesn't like appending. So, those np.hstack calls could be avoided for a slightly improved version as listed below -

def get_ranges_arr(starts,ends):
    counts = ends - starts
    counts_csum = counts.cumsum()
    id_arr = np.ones(counts_csum[-1],dtype=int)
    id_arr[0] = starts[0]
    id_arr[counts_csum[:-1]] = starts[1:] - ends[:-1] + 1
    return id_arr.cumsum()

Let's time it against our original approach -

In [151]: array1,array2 = (np.random.choice(range(1, 11),size=10**4, replace=True)\
     ...:                                      .cumsum().reshape(2, -1, order='F'))

In [152]: %timeit using_ones_cumsum(array1, array2)
1000 loops, best of 3: 276 µs per loop

In [153]: %timeit get_ranges_arr(array1, array2)
10000 loops, best of 3: 193 µs per loop

So, we have a 30% performance boost there!

Answer 3

Do you mean this?

In [440]: np.r_[10:14,65:70,200:204]
Out[440]: array([ 10,  11,  12,  13,  65,  66,  67,  68,  69, 200, 201, 202, 203])

or generalizing:

In [454]: np.r_[tuple([slice(i,j) for i,j in zip(array1,array2)])]
Out[454]: array([ 10,  11,  12,  13,  65,  66,  67,  68,  69, 200, 201, 202, 203])

Though this does involve a double loop, the explicit one to generate the slices and one inside r_ to convert the slices to arange.

    for k in range(len(key)):
        scalar = False
        if isinstance(key[k], slice):
            step = key[k].step
            start = key[k].start
                ...
                newobj = _nx.arange(start, stop, step)

I mention this because it shows that numpy developers consider your kind of iteration normal.

I expect that @unutbu's cleaver, if somewhat obtuse (I haven't figured out what it is doing yet), solution is your best chance of speed. cumsum is a good tool when you need to work with ranges than can vary in length. It probably gains most when the working with many small ranges. I don't think it works with overlapping ranges.

================

np.vectorize uses np.frompyfunc. So this iteration can also be expressed with:

In [467]: f=np.frompyfunc(lambda x,y: np.arange(x,y), 2,1)

In [468]: f(array1,array2)
Out[468]: 
array([array([10, 11, 12, 13]), array([65, 66, 67, 68, 69]),
       array([200, 201, 202, 203])], dtype=object)

In [469]: timeit np.concatenate(f(array1,array2))
100000 loops, best of 3: 17 µs per loop

In [470]: timeit np.r_[tuple([slice(i,j) for i,j in zip(array1,array2)])]
10000 loops, best of 3: 65.7 µs per loop

With @Darius's vectorize solution:

In [474]: timeit result = np.concatenate(ranges(array1, array2), axis=0)
10000 loops, best of 3: 52 µs per loop

vectorize must be doing some extra work to allow more powerful use of broadcasting. Relative speeds may shift if array1 is much larger.

@unutbu's solution isn't special with this small array1.

In [478]: timeit using_flatnonzero(array1,array2)
10000 loops, best of 3: 57.3 µs per loop

The OP solution, iterative without my r_ middle man is good

In [483]: timeit array3 = np.concatenate([np.arange(array1[i], array2[i]) for i in np.arange(0,len(array1))])
10000 loops, best of 3: 24.8 µs per loop

It's often the case that with a small number of loops, a list comprehension is faster than fancier numpy operations.

For @unutbu's larger test case, my timings are consistent with his - with a 17x speed up.

===================

For the small sample arrays, @Divakar's solution is slower, but for the large ones 3x faster than @unutbu's. So it has more of a setup cost, but scales slower.

Answer 4

This is my approach combining vectorize and concatenate:

Implementation:

import numpy as np

array1, array2 = np.array([10, 65, 200]), np.array([14, 70, 204])

ranges = np.vectorize(lambda a, b: np.arange(a, b), otypes=[np.ndarray])
result = np.concatenate(ranges(array1, array2), axis=0)

print result
# [ 10  11  12  13  65  66  67  68  69 200 201 202 203]

---

Performance:

%timeit np.concatenate(ranges(array1, array2), axis=0)

100000 loops, best of 3: 13.9 µs per loop