python - 计算系列开始/结束处的零数

Question

我有这样一个系列:

s = pd.Series([0, 0, 0, 1, 2, 3])
s
Out[00]: 
0    0
1    0
2    0
3    1
4    2
5    0
dtype: int64

我想计算这个系列中开始和尾随零的数量。所以在这种情况下，我应该得到 3 作为开始，因为在第一个非零数字之前有 3 个零，而尾随零为 1，因为在最后一个非零之后的系列尾部有一个零。

到目前为止我做了什么

到目前为止我的解决方案是使用累积和

sum(s.cumsum() == 0) # begenning
np.sum(np.cumsum(s.values[::-1]) == 0) # trailing

但这对于非常大的系列来说非常慢，尤其是尾随零计算，我需要一个替代方案。

Answer 1

使用numpy.nonzero:

import numpy as np

n_rows = len(s)
indices = np.nonzero(s)[0]

if indices.size>0:
    head = indices[0]
    trail = n_rows - indices[-1] -1
else:
    head, trail = n_rows, n_rows
print(head, trail)

输出:

3 1

基准测试(大约快 15 倍):

s = np.zeros(100000)
s[30000:50000] +=1
s = pd.Series(s)

%%timeit

n_rows = len(s)
indices = np.nonzero(s)[0]

if indices.size>0:
    head = indices[0]
    trail = n_rows - indices[-1] -1
else:
    head, trail = n_rows, n_rows
# 661 µs ± 8.63 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%%timeit

sum(s.cumsum() == 0) # begenning
np.sum(np.cumsum(s.values[::-1]) == 0) # trailing
# 9.39 ms ± 163 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

测试和编辑:适用于全零和非零情况。