python的部分累积和

Question

假设我有一个numpy数组(或者pandas Series，如果它使它变得更简单)，它如下所示：

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

我想要转换成一个数组

bar = np.array([0, 1, 2, 3, 4,0, 1, 2, 0, 1, 2, 3])

其中的条目是您需要向左走多少步才能在foo中找到一个foo。

现在，显然可以编写一个循环来从bar中计算foo，但这将是缓慢的。还有什么更聪明的人可以做的吗？

更新 pd.Series解决方案比纯numpy解决方案慢7倍。愚蠢的循环解决方案非常慢(不足为奇)，但是当用numba编译jit时，速度与numpy解决方案一样快。

Answer 1

你可以用cumcount做pandas

s = pd.Series(foo)
bar = s.groupby(s.cumsum()).cumcount().to_numpy()
Out[13]: array([0, 1, 2, 3, 4, 0, 1, 2, 0, 1, 2, 3], dtype=int64)

Answer 2

有一个选项，特别是对于共享示例，使用numpy：

# get positions where value is 1
pos = foo.nonzero()[0]
# need this when computing the cumsum
values = np.diff(pos) - 1
arr = np.ones(foo.size, dtype=int)
arr[0] = 0
arr[pos[1:]] = -values
arr.cumsum()

array([0, 1, 2, 3, 4, 0, 1, 2, 0, 1, 2, 3])

Answer 3

是的，您可以使用numpy.cumsum()函数来非常有效地计算foo中的bar。以下是你如何做到的：

import numpy as np

# Define the input array
foo = np.array([1, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0])

# Compute the cumulative sum of foo, but shift the values
# to the right by one position and insert a 0 at the beginning
cs = np.insert(np.cumsum(foo), 0, 0)

# Subtract the shifted cumulative sum from the original cumulative sum
bar = np.cumsum(foo) - cs

# Print the result
print(bar)

这将为您提供所需的输出：array([0, 1, 2, 3, 4, 0, 1, 2, 0, 1, 2, 3])。

下面简要解释一下这是如何工作的：numpy.cumsum()函数计算数组的累积和，这意味着它将数组的所有元素加到并包含每个元素。因此，例如，如果将numpy.cumsum()应用于数组[1, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0]，则得到数组[1, 1, 1, 1, 1, 2, 2, 2, 3, 3, 3, 3]。

接下来，我们使用numpy.insert()函数将累积和的值右移一个位置，并在开头插入0。这就给出了数组[0, 1, 1, 1, 1, 1, 2, 2, 2, 3, 3, 3]。

最后，从原始累积和中减去移位累积和，从而得到阵列[0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0]。这与原始的数组foo完全相同，只不过1s已经被向左走去寻找下一个1所需的步骤数所取代，这正是我们想要的结果。

我希望这能帮到你!如果你还有其他问题，请告诉我。