Найти число нулей до ненулевого в массиве numpy

Вот итеративная версия Cython, которая может быть вашим лучшим выбором, если это серьезное узкое место

 # saved as file count_leading_zeros.pyx import numpy as np cimport numpy as np cimport cython DTYPE = np.int ctypedef np.int_t DTYPE_t @cython.boundscheck(False) def count_leading_zeros(np.ndarray[DTYPE_t, ndim=1] a): cdef int elements = a.size cdef int i = 0 cdef int count = 0 while i < elements: if a[i] == 0: count += 1 else: return count i += 1 return count 

Это похоже на ответ @ mtrw, но с индексированием на собственных скоростях. Мой Cython немного отрывочен, поэтому могут быть сделаны дополнительные улучшения.

Быстрый тест чрезвычайно благоприятного случая с IPython с несколькими различными методами

 In [1]: import numpy as np In [2]: import pyximport; pyximport.install() Out[2]: (None, <pyximport.pyximport.PyxImporter at 0x53e9250>) In [3]: import count_leading_zeros In [4]: %paste def count_leading_zeros_python(x): ctr = 0 for k in x: if k == 0: ctr += 1 else: return ctr return ctr ## -- End pasted text -- In [5]: a = np.zeros((10000000,), dtype=np.int) In [6]: a[5] = 1 In [7]: In [7]: %timeit np.min(np.nonzero(np.hstack((a, 1)))) 10 loops, best of 3: 91.1 ms per loop In [8]: In [8]: %timeit np.where(a)[0][0] if np.shape(np.where(a)[0])[0] != 0 else np.shape(a)[0] 10 loops, best of 3: 107 ms per loop In [9]: In [9]: %timeit count_leading_zeros_python(a) 100000 loops, best of 3: 3.87 µs per loop In [10]: In [10]: %timeit count_leading_zeros.count_leading_zeros(a) 1000000 loops, best of 3: 489 ns per loop 

Однако я бы использовал только что-то подобное, если бы у меня были доказательства (с профайлером), что это было узким местом. Многие вещи могут показаться неэффективными, но никогда не стоить вашего времени, чтобы исправить.

Добавив ненулевой номер в конце массива, вы все равно можете использовать np.nonzero для получения желаемого результата.

 A = np.array([0,1,2]) B = np.array([0,0,0]) np.min(np.nonzero(np.hstack((A, 1)))) # --> 1 np.min(np.nonzero(np.hstack((B, 1)))) # --> 3 

 i = np.argmax(A!=0) if i==0 and np.all(A==0): i=len(A) 

Это должно быть наиболее эффективным решением без расширений. Также легко векторизован, чтобы действовать по нескольким осям.

Что не так с наивным подходом:

 def countLeadingZeros(x): """ Count number of elements up to the first non-zero element, return that count """ ctr = 0 for k in x: if k == 0: ctr += 1 else: #short circuit evaluation, we found a non-zero so return immediately return ctr return ctr #we get here in the case that x was all zeros 

Это возвращает, как только будет найден ненулевой элемент, поэтому это O (n) в худшем случае. Вы можете сделать это быстрее, поместив его на C, но было бы целесообразно проверить, действительно ли это необходимо для массивов, с которыми вы работаете.

Я удивлен, почему никто не использовал np.where еще

np.where(a)[0][0] if np.shape(np.where(a)[0])[0] != 0 else np.shape(a)[0] сделает трюк

 >> a = np.array([0,1,2]) >> np.where(a)[0][0] if np.shape(np.where(a)[0])[0] != 0 else np.shape(a)[0] ... 1 >> a = np.array([0,0,0)) >> np.where(a)[0][0] if np.shape(np.where(a)[0])[0] != 0 else np.shape(a)[0] ... 3 >> a = np.array([1,2,3)) >> np.where(a)[0][0] if np.shape(np.where(a)[0])[0] != 0 else np.shape(a)[0] ... 0 

Если вас не волнует скорость, у меня есть небольшой трюк, чтобы выполнить эту работу:

 a = np.array([0,0,1,1,1]) t = np.where(a==0,1,0)+np.append(np.where(a==0,0,1),0)[1:] print t [1 2 1 1 0] np.where(t==2) (array([1]),)