python - 如何计算到 Pandas 系列中前一个零的距离？

Question

我有以下 pandas 系列(表示为列表):

[7,2,0,3,4,2,5,0,3,4]

我想定义一个新系列，将距离返回到最后一个零。这意味着我想要以下输出:

[1,2,0,1,2,3,4,0,1,2]

如何以最有效的方式在 pandas 中做到这一点？

Answer 1

复杂度为O(n)。什么会减慢它在 python 中执行 for 循环。如果序列中有 k 个零，并且 log k 与序列的长度相比可以忽略不计，O(n log k) 解决方案会是:

>>> izero = np.r_[-1, (ts == 0).nonzero()[0]]  # indices of zeros
>>> idx = np.arange(len(ts))
>>> idx - izero[np.searchsorted(izero - 1, idx) - 1]
array([1, 2, 0, 1, 2, 3, 4, 0, 1, 2])