在 1 和 0 的一维 numpy 数组中，如何将 1 之后的下 n 个元素转换为 0？

Question

对于 1 和 0 的一维 numpy 数组，我如何有效地“屏蔽”数组，以便在出现 1 后，将数组的下 n 个元素转换为零。在 n 个元素通过后，模式会重复，这样会保留下一次出现的第一个 1，然后再保留 n 个零。

保留第一个符合条件的 1 很重要，因此可以使用简单的掩码，例如： [true, false, false, true ...] 行不通。 此外，数据集非常庞大，因此效率很重要。

我已经编写了粗略的 python 代码来给我想要的结果，但它对于我需要的东西来说太慢了。 这是一个例子：

data = [0, 0, 1, 1, 0, 1, 1, 1, 0, 0, 0, 1]
n = 3

newData = []
tail = 0
for x in data:
    if x == 1 and tail <= 0:
        newData.append(1)
        tail = n
    else:
        newData.append(0)
        tail -= 1
print(newData)

新数据：[0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 1]

这个问题有可能有一个矢量化的 numpy 解决方案吗？ 我正在处理数以万计的数组，每个数组中有超过一百万个元素。到目前为止，使用 numpy 函数是管理此问题的唯一方法。

Answer 1

你可以这样做：-

N = 3
data = [0, 0, 1, 1, 0, 1, 1, 1, 0, 0, 0, 1, 0, 1, 1]
newData = data.copy()
i = 0
M = [0 for _ in range(N)]
while i < len(newData) - N:
    if newData[i] == 1:
        newData[i + 1:i + 1 + N] = M
        i += N
    i += 1

print(newData)

Answer 2

据我所知，在 numpy 中没有完全可以做到这一点的选项。不过，您仍然可以使用 numpy 来减少获取索引的时间。

data = [0, 0, 1, 1, 0, 1, 1, 1, 0, 0, 0, 1]
n=3
def get_new_data(data,n):
  new_data = np.zeros(len(data))
  non_zero = np.argwhere(data).ravel()
  idx = non_zero[0]
  new_data[idx] =1
  idx += n
  for i in non_zero[1:]:
    if i > idx:
      new_data[i] = 1
      idx+=n
  return new_data
get_new_data(data, n)

这样的函数应该会给你一个更好的运行时间，因为你没有循环整个数组。

如果这对你来说仍然不是最优的，你可以看看使用numba，它与 numpy 配合得很好，并且相对容易使用。