Numpy: Using two different sets of boolean masks at once using advanced indexing-CodePudding

Say that I have this array

import numpy as np
a = np.array(
    [
     [[1,2], [3,4], [5,6]],
     [[11,21], [31,41], [51,61]],
     [[12,22], [32,42], [52,62]],
    ]
)

And these two boolean masks

starts = array([[False,  True, False],
       [ True, False, False],
       [ True, False, False]])

ends = array([[False, False,  True],
       [False,  True, False],
       [False, False,  True]])

I would like to use both of them to index a so that each row is combined. This is the desired result

array([[ 3,  4,  5,  6],
       [11, 21, 31, 41],
       [12, 22, 52, 62]])

Which I can get through np.concatenate((a[starts], a[ends]), axis=1)

But I have heard that concatenate is not computationally efficient compared to advanced indexing. I have a feeling that it might be possible to get my desired result through advanced indexing reshaping since I am indexing the same array twice, but I am unable to come up with a solution.

Is there a more computationally more efficient way to get my desired result than np.concatenate((a[starts], a[ends]), axis=1)?

CodePudding user response：

The best I can do using advanced indexing and reshape is

def foo(a,starts,ends):
    I,J = [np.stack([i,j], axis=1) for i,j in zip(np.nonzero(starts),np.nonzero(ends))]
    res = a[I,J]
    return res.reshape(3,4)

In [123]: timeit np.concatenate((a[starts], a[ends]), axis=1)
16.3 µs ± 14.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
In [124]: timeit foo(a,starts,ends)
42.8 µs ± 105 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In other SO people have asked about combining multiple slices, worrying about the "iteration" and concatenation. But the alternative is to expand the slices into advanced indexing, and doing one just one indexing operation. Generally I've those alternatives to be competitive.

Simple sliced indexing is fast because creates a view - a new array with its own shape and strides, but shared memory (data-buffer). Advanced indexing makes a copy - which may take time if the array is large. Concatenate also involves a copy - all elements of the sources arrays have to be copied to the result. What you don't want to do is copy those elements again and again.