Python Retain函数。在计算中使用上一行的值

Question

In [10]: df
Out[10]:
     PART AVAILABLE_INVENTORY DEMAND
1    A    12                  6
2    A    12                  2
3    A    12                  1
4    B    24                  1
5    B    24                  1
6    B    24                  4
7    B    24                  3

想要的输出：

     PART AVAILABLE_INVENTORY DEMAND  AI   AI_AFTER
1    A    12                  6       12   6
2    A    12                  2       6    4
3    A    12                  1       4    3
4    B    24                  1       24   23
5    B    24                  1       23   22
6    B    24                  4       22   18
7    B    24                  3       18   15

到目前为止，我的代码在下面，但是没有提供我想要的输出：

def retain(df):
    df['PREV_PART'] = df['PART'].shift()
    df['PREV_AI_AFTER'] = df['AI'].shift() - df['DEMAND'].shift()
    df['AI'] = np.where(df['PART'] != df['PREV_PART'], df['AI'], df['PREV_AI_AFTER'])
    df['AI_AFTER'] = df['AI'] - df['DEMAND']

df['AI'] = df['AVAILABLE_INVENTORY']
retain(df)

考虑性能的最快方法是什么？

Answer 1

您可以通过之前创建的[DEMAND]列中的groupby和[AI_AFTER]列中的cumsum来使用shift进行此操作：

df['AI_AFTER'] = df['AVAILABLE_INVENTORY'] - df.groupby('PART')['DEMAND'].cumsum()
df['AI'] = df.groupby('PART')['AI_AFTER'].shift().fillna(df['AVAILABLE_INVENTORY'])
print (df)
  PART  AVAILABLE_INVENTORY  DEMAND  AI_AFTER    AI
1    A                   12       6         6  12.0
2    A                   12       2         4   6.0
3    A                   12       1         3   4.0
4    B                   24       1        23  24.0
5    B                   24       1        22  23.0
6    B                   24       4        18  22.0
7    B                   24       3        15  18.0

Answer 2

VERRRY与Ben.T's Answer相似。如果您喜欢这种方法，请选择他们的答案。这就是我安排流程的方式。

def f(d):
    i = d.AVAILABLE_INVENTORY
    c = d.DEMAND.cumsum()
    return pd.concat({'AI': i - c.shift(fill_value=0), 'AI_AFTER': i - c}, axis=1)

df.join(df.groupby('PART').apply(f))

  PART  AVAILABLE_INVENTORY  DEMAND  AI  AI_AFTER
1    A                   12       6  12         6
2    A                   12       2   6         4
3    A                   12       1   4         3
4    B                   24       1  24        23
5    B                   24       1  23        22
6    B                   24       4  22        18
7    B                   24       3  18        15