我有一个数据框架（DF），其中按天汇总存款和取款，我想知道计算每天余额的最快方法是什么。因为它必须能够扩展。欢迎使用 Pandas 和 Spark 提供答案！以下是输入 DF 的示例：
输入
| 日期 | 存款 | 提款 |
|---|---|---|
| 2024-01-01 | 100.00 | 0.00 |
| 2024-01-02 | 0.00 | 0.00 |
| 2024-01-03 | 50.00 | 30.00 |
| 2024-01-04 | 0.00 | 0.00 |
| 2024-01-05 | 0.00 | 200.00 |
| 2024-01-06 | 20.00 | 0.00 |
| 2024-01-07 | 20.00 | 0.00 |
| 2024-01-08 | 0.00 | 0.00 |
这些存款和取款来自每天收益率为 10% 的投资账户。除非余额为负数。在这种情况下，每日收益必须为零。获取
daily_returnbalance动向 = 前一天余额 + 存款 - 提款
如果变动 > 0，则利息 = 0.1，否则 0
每日收益 = 动量 * 利息
余额=动量+每日收益
下面是所需输出 DF 的示例：
所需输出
| 日期 | 存款 | 提款 | 每日回报 | 平衡 |
|---|---|---|---|---|
| 2024-01-01 | 100.00 | 0.00 | 10.00 | 110.00 |
| 2024-01-02 | 0.00 | 0.00 | 11.00 | 121.00 |
| 2024-01-03 | 50.00 | 30.00 | 10.14 | 155.10 |
| 2024-01-04 | 0.00 | 0.00 | 15.51 | 170.61 |
| 2024-01-05 | 0.00 | 200.00 | 0.00 | -29.39 |
| 2024-01-06 | 20.00 | 0.00 | 0.00 | -9.39 |
| 2024-01-07 | 20.00 | 0.00 | 1.06 | 11.67 |
| 2024-01-08 | 0.00 | 0.00 | 1.17 | 12.84 |
我有什么
我在 Pandas 中有一个解决方案可以实现所需的输出，但是它会迭代 DF 的每一行，即速度很慢。有没有办法向量化这个计算以加快速度？或者也许另一种方法？这是我的实现：
import pandas as pd
df = pd.DataFrame({
"date": pd.date_range(start="2024-01-01", end="2024-01-08"),
"deposit": [100.0, 0.0, 50.0, 0.0, 0.0, 20.0, 20.0, 0.0],
"withdrawal": [0.0, 0.0, 30.0, 0.0, 200.0, 0.0, 0.0, 0.0]
})
daily_returns = []
balances = []
prev_balance = 0
for _, row in df.iterrows():
movements = prev_balance + row["deposit"] - row["withdrawal"]
interest = 0.1 if movements > 0 else 0
daily_return = movements * interest
balance = movements + daily_return
daily_returns.append(daily_return)
balances.append(balance)
prev_balance = balance
df["daily_return"] = daily_returns
df["balance"] = balances
对于这种类型的计算，我会使用 numba，例如：
from numba import njit
@njit
def calculate(deposits, withdrawals, out_daily_return, out_balance):
prev_balance = 0
for i, (deposit, withdrawal) in enumerate(zip(deposits, withdrawals)):
movements = prev_balance + deposit - withdrawal
interest = 0.1 if movements > 0 else 0
daily_return = movements * interest
balance = movements + daily_return
out_daily_return[i] = daily_return
out_balance[i] = balance
prev_balance = balance
df["daily_return"] = 0.0
df["balance"] = 0.0
calculate(
df["deposit"].values,
df["withdrawal"].values,
df["daily_return"].values,
df["balance"].values,
)
print(df)
打印：
date deposit withdrawal daily_return balance
0 2024-01-01 100.0 0.0 10.0000 110.0000
1 2024-01-02 0.0 0.0 11.0000 121.0000
2 2024-01-03 50.0 30.0 14.1000 155.1000
3 2024-01-04 0.0 0.0 15.5100 170.6100
4 2024-01-05 0.0 200.0 -0.0000 -29.3900
5 2024-01-06 20.0 0.0 -0.0000 -9.3900
6 2024-01-07 20.0 0.0 1.0610 11.6710
7 2024-01-08 0.0 0.0 1.1671 12.8381