python - 一维numpy数组中到下一个非零元素的距离

Question

我有一个由 1 和 0 组成的一维 numpy 数组，如下所示：

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

对于数组的每个非零元素，我想计算到下一个非零元素的“距离”。也就是说，我想回答“下一个非零元素还有多远？”这个问题。所以上述数组的结果是：

[0, 0, 0, 1, 3, 0, 0, 6, 0, 0, 0, 0, 0, 4, 0, 0, 0, 0]

是否有内置的 numpy 函数？如果没有，在 numpy 中实现这一点的最有效方法是什么？

标签： pythonnumpy

Answer 1

咱们试试吧：

import numpy as np

arr = np.array([0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1])

# create output
res = np.zeros_like(arr)

# select indices non-zero
where, = np.where(arr)

# assign the indices of the non-zero the diff
res[where[:-1]] = np.diff(where)
print(res)

输出

[0 0 0 1 3 0 0 6 0 0 0 0 0 4 0 0 0 0]