numpy 如何计算每次遇到零时重置的1的累积和

ijxebb2r 于 2023-01-26 发布在其他

关注(0)|答案(2)|浏览(143)

我有一个由0和1组成的数组。我想计算所有连续1的累积和，每次遇到0时重置，使用numpy，因为我有数千个包含数千行和列的数组。
我可以用循环来做，但我怀疑它不会有效。你有一个更聪明和快速的方法来运行它的数组吗？下面是输入和预期输出的简短示例：

import numpy as np
arr_in = np.array([[1,1,1,1,1,1], [0,0,0,0,0,0], [1,0,1,0,1,1], [0,1,1,1,0,0]])
print(arr_in)
print("expected result:")
arr_out = np.array([[1,2,3,4,5,6], [0,0,0,0,0,0], [1,0,1,0,1,2], [0,1,2,3,0,0]])
print(arr_out)

运行时：

[[1 1 1 1 1 1]
 [0 0 0 0 0 0]
 [1 0 1 0 1 1]
 [0 1 1 1 0 0]]
expected result:
[[1 2 3 4 5 6]
 [0 0 0 0 0 0]
 [1 0 1 0 1 2]
 [0 1 2 3 0 0]]

numpy

来源：https://stackoverflow.com/questions/75233896/how-to-calculate-cumulative-sums-of-ones-with-a-reset-each-time-a-zero-is-encoun

2条答案

按热度按时间

nbnkbykc1#

可以计算1的cumsum，然后识别0并向前填充累积和以减去它：

# identify 0s
mask = arr_in==0

# get classical cumsum
cs = arr_in.cumsum(axis=1)

# ffill the cumsum value on 1s
# subtract from cumsum
out = cs-np.maximum.accumulate(np.where(mask, cs, 0), axis=1)

输出：

[[1 2 3 4 5 6]
 [0 0 0 0 0 0]
 [1 0 1 0 1 2]
 [0 1 2 3 0 0]]

第二个示例的输出：

[[1 2 3 4 5 6 0 1]
 [0 1 2 0 0 0 1 0]]

赞(0）回复(0）举报 2023-01-26

7cjasjjr2#

使用numba.vectorize，您可以定义一个定制的numpy ufunc用于累加。

import numba as nb # v0.56.4, no support for numpy >= 1.22.0
import numpy as np # v1.21.6

@nb.vectorize([nb.int64(nb.int64, nb.int64)])
def reset_cumsum(x, y):
    return x + y if y else 0

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

reset_cumsum.accumulate(arr_in, axis=1)

产出

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

赞(0）回复(0）举报 2023-01-26

我来回答

numpy 如何计算每次遇到零时重置的1的累积和

2条答案

相关问题

热门标签

最新问答