[Numpy-discussion] index of the first element fulfilling a condition?

Erin Sheldon erin.sheldon at gmail.com
Thu Aug 19 21:05:50 EDT 2010


Excerpts from Johann Bauer's message of Wed Aug 18 14:07:48 -0400 2010:
> Hi,
> 
> is there a good method to find the index of the first element in a 1D 
> array fulfilling a condition?
> 
> The following does the job
> 
> >>> import numpy
> >>> a = numpy.array([1,5,78,3,12,4])
> >>> numpy.where( a>10 )[0][0]
> 2
> 
> but it first compares the entire array and then selects the first index. 
> Is there a way (especially for big arrays) to stop the comparison after 
> the first hit?
> 
> Thanks for any help, Johann

If you are willing to use scipy.weave, it is straightforward:

import numpy                                                                             
from scipy import weave                                                                  
                                                                                         
def get_first_greater(arr, value):                                                       
    """                                                                                  
    Return index of first element in array that is greater than value.                   
    """                                                                                  
                                                                                         
    code = """                                                                           
    int64_t i=0, index=-1;                                                               
    while (i < arr.size()) {                                                             
        if ( arr(i) > value ) {                                                          
            index = i;                                                                   
            break;                                                                       
        }                                                                                
        i++;                                                                             
    }                                                                                    
    return_val = index;                                                                  
    """                                                                                  
                                                                                         
    index = weave.inline(code, ['arr','value'],                                          
                         type_converters = weave.converters.blitz)                       
    return index 


Caveats: In scipy 0.7 and gcc >= 4.3 this probably won't compile.  It
does work for scipy 0.8.

Erin Scott Sheldon



More information about the NumPy-Discussion mailing list