时间戳

首页 > 解决方案 > 在 1 和 0 的一维 numpy 数组中，如何将 1 之后的下 n 个元素转换为 0？

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

问题描述

对于 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 函数是管理此问题的唯一方法。

标签： pythonnumpyvectorizationnumpy-ndarraymasked-array

解决方案

据我所知，在 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 配合得非常好，并且相对易于使用。

推荐阅读