如果超出限制，则重置 cumsum（python）

Question

以下 numpy sn-p 将返回输入数组的 cumsum，每次遇到 NaN 时都会重置。

v = np.array([1., 1., 1., np.nan, 1., 1., 1., 1., np.nan, 1.])
n = np.isnan(v)
a = ~n
c = np.cumsum(a)
d = np.diff(np.concatenate(([0.], c[n])))
v[n] = -d
result = np.cumsum(v)

以类似的方式，如果使用矢量化 pandas 或 numpy 操作，如果 cumsum 超过某个值，我如何计算重置的 cumsum？

例如对于限制 = 5，输入 = [1,1,1,1,1,1,1,1,1,1]，输出 = [1,2,3,4,5,1,2,3,4, 5]

Answer 1

如果您的数组中的数字都是正数，那么使用cumsum() 然后使用模运算符可能是最简单的：

>>> a = np.array([1,1,1,1,1,1,1,1,1,1])
>>> limit = 5
>>> x = a.cumsum() % limit
>>> x
array([1, 2, 3, 4, 0, 1, 2, 3, 4, 0])

然后您可以将任何零值设置回限制以获取所需的数组：

>>> x[x == 0] = limit
>>> x
array([1, 2, 3, 4, 5, 1, 2, 3, 4, 5])

---

这是使用 Pandas 的 expanding_apply 方法的一种可能的通用解决方案。 （我没有对它进行广泛的测试......）

先定义一个修改后的cumsum函数：

import pandas as pd

def cumsum_limit(x):
    q = np.sum(x[:-1])
    if q > 0:
        q = q%5
    r = x[-1]
    if q+r <= 5:
        return q+r
    elif (q+r)%5 == 0:
        return 5
    else:
        return (q+r)%5

a = np.array([1,1,1,1,1,1,1,1,1,1]) # your example array

像这样将函数应用于数组：

>>> pd.expanding_apply(a, lambda x: cumsum_limit(x))
array([ 1.,  2.,  3.,  4.,  5.,  1.,  2.,  3.,  4.,  5.])

这是另一个更有趣的系列的函数：

>>> s = pd.Series([3, -8, 4, 5, -3, 501, 7, -100, 98, 3])
>>> pd.expanding_apply(s, lambda x: cumsum_limit(x)) 
0     3
1    -5
2    -1
3     4
4     1
5     2
6     4
7   -96
8     2
9     5
dtype: float64