python - 累积总和和结转 - 用 Pandas 矢量化

Question

我可以在 a 和 b 上使用 pandas.Series 数学运算而不显式循环吗？

In [38]: a = pd.Series([4, 8, 3, 6, 2])

In [39]: b = pd.Series([3, 9, 5, 5, 4])

In [40]: alist = a.tolist()
    ...: blist = b.tolist()
    ...: for i in range(len(alist)):
    ...:     diff = max(0, alist[i] - blist[i])
    ...:     try:
    ...:         alist[i + 1] = alist[i + 1] + diff
    ...:     except IndexError:
    ...:         if diff > 0:
    ...:             alist.append(diff)
    ...:     blist[i] = max(0, blist[i] - alist[i])
    ...: 

In [41]: alist
Out[41]: [4, 9, 3, 6, 3]

In [42]: blist
Out[42]: [0, 0, 2, 0, 1]

如果 a 和 b 的差值大于零，我将递增 a 的下一个值，然后从类似累积和的计算中减去 b。

Answer 1

这是使用 numpy 的一种方式:

import numpy as np

a += np.maximum(0, a-b).shift().fillna(0).astype(int)
b = np.maximum(0, b - a)

print(a)

0    4
1    9
2    3
3    6
4    3
dtype: int64

print(b)

0    0
1    0
2    2
3    0
4    1
dtype: int64