如何根据期末余额和未来预测动态计算周供应量？

Question

我有以下 2 个表/数据框：

data1 = {
    'group': ['A', 'A', 'B'],
    'category': [1, 2, 2],
    'start_balance': [300, 3000, 80]
}

df1 = pd.DataFrame(data1)
df1

data2 = {
    'group': ['A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'B', 'B', 'B', 'B', 'B'],
    'category': [1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
    'week': [1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5],
    'forecast': [82, 273, 232, 176, 142, 972, 1041, 503, 510, 566, 19, 15, 15, 13, 14],
    'end_balance': [400, 520, 288, 395, 253, 188, -853, -1356, -1866, -2432, 61, 40, 80, 67, 53]
}

df2 = pd.DataFrame(data2)
df2

df1

代表组/类别列表（键）和起始余额。

df2

显示每个键的预测和期末余额是多少。请注意，最终余额是多种因素的结果，此处未包含这些因素，因为它们与任务无关。

我需要完成：

1. 向

   "wks_of_supply"

   添加一列

   df1

   ，根据

   "start_balance"

   中的

   df1

   和

   "forecast"

   中相同键的

   df2

   计算剩余周数。从第一周开始，需要确定可以全额承保多少周，然后按比例分配最后一周。如果余额一直覆盖到预测结束将导致 inf（注意仍然需要是浮点数）。

这是预期的结果：

组	类别	开始余额	供应周数
A	1	300	1.798534799
A	2	3000	3.949019608
B	2	80	inf

2. 对

   df2

   应用类似的逻辑来添加

   "wks_of_supply"

   列。请注意，预测将从第 2 周开始，因为第 1 周预测已用于

   "start_balance"

   中的

   "wks_of_supply"

   df1

   。因此，最后一周将永远是

   NaN

   。还需要考虑负数和

   -inf

   "wks_of_supply"

   。

这是预期的结果：

组	类别	周	预测	结束余额	供应周数
A	1	1	82	400	1.547413793
A	1	2	273	520	2.788732394
A	1	3	232	288	1.788732394
A	1	4	176	395	inf
A	1	5	142	253	NaN
A	2	1	972	188	0.180595581
A	2	2	1041	-853	-1.68627451
A	2	3	503	-1356	-inf
A	2	4	510	-1866	-inf
A	2	5	566	-2432	NaN
B	2	1	19	61	inf
B	2	2	15	40	2.857142857
B	2	3	15	80	inf
B	2	4	13	67	inf
B	2	5	14	53	NaN

需要在大约 200 万行数据集上运行它，因此需要尽可能快地运行的东西，尽管可能很难矢量化。

尝试创建一个键列表，然后运行

for

循环来查找余额变为负值的那一周。然后将结果读入 df 并计算部分周分数。这只能解决第一个问题，而且速度非常慢。

Answer 1

使用矢量化代码是可能的，尽管在一个地方我发现我需要从 Pandas 中删除到 Numpy 来进行索引：

import numpy as np
import pandas as pd


df1 = pd.DataFrame({
    'group':         ( 'A',  'A', 'B'),
    'category':      (  1,    2,   2),
    'start_balance': (300, 3000,  80),
}).set_index(['group', 'category'])

df2 = pd.DataFrame({
    'group': ('A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'B', 'B', 'B', 'B', 'B'),
    'category': (1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2),
    'week': (1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5),
    'forecast': (82, 273, 232, 176, 142, 972, 1041, 503, 510, 566, 19, 15, 15, 13, 14),
    'end_balance': (400, 520, 288, 395, 253, 188, -853, -1356, -1866, -2432, 61, 40, 80, 67, 53),
}).set_index(['group', 'category', 'week'])

declining_balance = (
    df1['start_balance'] -
    df2['forecast'].groupby(['group', 'category']).cumsum()
).unstack('week').iloc[:, ::-1]

exhausted_idx = declining_balance.apply(pd.Series.searchsorted, axis=1, value=0)
is_exhausted = declining_balance.iloc[:, 0] <= 0
whole_weeks = declining_balance.shape[1] - exhausted_idx
whole_weeks[~is_exhausted] = np.inf
y_idx = np.arange(declining_balance.shape[0])[is_exhausted]

numerator = declining_balance.values[
    y_idx,
    exhausted_idx[is_exhausted],
]
denominator = df2.forecast.unstack('week').values[
    y_idx,
    (5 - exhausted_idx)[is_exhausted],
]

weeks_of_supply = whole_weeks.copy()
weeks_of_supply[is_exhausted] += numerator/denominator
print(weeks_of_supply)

group  category
A      1           1.798535
       2           3.949020
B      2                inf
dtype: float64