Numpy equivalent of list.index

前端 未结 6 2260
粉色の甜心
粉色の甜心 2021-02-12 11:48

In a low-level function that is called many times, I need to do the equivalent of python\'s list.index, but with a numpy array. The function needs to return when it finds the f

6条回答
  •  傲寒
    傲寒 (楼主)
    2021-02-12 12:07

    You can code it in Cython and just import from a Python script. There is no need to migrate your entire project into Cython.

    # paste into: indexing.pyx
    def index(long[:] lst, long value):
        cdef int i
        for i in range(len(lst)):
            if lst[i] == value:
                return i
        raise ValueError
    
    # import in your .py code
    import pyximport
    pyximport.install()
    from indexing import index
    
    # example
    from numpy import zeros
    a = zeros(10**6, int)
    a[-1] = 1
    
    index(a, 1)
    Wall time: 6.07 ms
    999999
    
    index(a, 0)
    Wall time: 38.1 µs
    0
    

提交回复
热议问题