python - 如何计算数组中相邻的重复元素？

Question

我有一个由 0 和 1 组成的数组

[0,0,1,1,1,0,0,0,0,1,1,0,0]

我想定义一个函数，它将此数组作为输入并输出一个长度相同的数组，其中包含第一个 1 出现的索引中相邻 1 的计数(否则为 0)。所以输出将是

[0,0,3,0,0,0,0,0,0,2,0,0,0]

因为 1 连续 3 次出现在第 2 个索引中，而 1 连续 2 次出现在第 9 个索引中。

有没有办法使用 numpy 来做到这一点？如果没有，是否有一些(有效的)pythonic 方法来做到这一点？

Answer 1

这是一个使用纯向量化操作且没有列表迭代的解决方案:

import numpy as np

data = np.array([0,0,1,1,1,0,0,0,0,1,1,0,0])
output = np.zeros_like(data)

where = np.where(np.diff(data))[0]
vals = where[1::2] - where[::2]
idx = where[::2] + 1

output[idx] = vals
output
# array([0, 0, 3, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0])