python - DataFrame 计算平均购买价格

Question

我有一个包含两列的数据框:数量和价格。

df = pd.DataFrame([
[ 1, 5],
[-1, 6],
[ 2, 3],
[-1, 2],
[-1, 4],
[ 1, 2],
[ 1, 3],
[ 1, 4],
[-2, 5]], columns=['quantity', 'price'])

df['amount'] = df['quantity'] * df['price']
df['cum_qty'] = df['quantity'].cumsum()

我添加了两个新列 amount 和 cum_qty(累计数量)。 现在数据框看起来像这样(正数代表买入，负数代表卖出):

   quantity  price  amount  cum_qty
0         1      5       5        1
1        -1      6      -6        0
2         2      3       6        2
3        -1      2      -2        1
4        -1      4      -4        0
5         1      2       2        1
6         1      3       3        2
7         1      4       4        3
8        -2      5     -10        1

我想计算平均买入价。

每次当 cum_qty = 0 时，数量和金额都应重置为零。 所以我们正在查看索引 = [5,6,7] 的行。 对于每一行，一件商品以价格 2、3 和 4 购买，这意味着我有 3 件商品，平均价格为 3 [(2 + 3 + 4)/3]。

在索引 = 8 的卖出发生后(卖出交易不会改变买入价格)，我将以 3 的价格各持有一个。

因此，基本上，我必须将所有累计购买金额除以累计数量(上次累计数量不为零)。

如何使用 pandas DataFrame 计算所有交易的结果？

Answer 1

这是一个使用循环的不同解决方案:

import pandas as pd
import numpy as np

# Original data
df = pd.DataFrame({
    'quantity': [ 1, -1,  2, -1, -1,  1,  1,  1, -2],
    'price': [5, 6, 3, 2, 4, 2, 3, 4, 5]
})

# Process the data and add the new columns
df['amount'] = df['quantity'] * df['price']
df['cum_qty'] = df['quantity'].cumsum()
df['prev_cum_qty'] = df['cum_qty'].shift(1, fill_value=0)
df['average_price'] = np.nan
for i, row in df.iterrows():
    if row['quantity'] > 0:
        df.iloc[i, df.columns == 'average_price' ] = (
            row['amount'] +
            df['average_price'].shift(1, fill_value=df['price'][0])[i] *
            df['prev_cum_qty'][i]
        )/df['cum_qty'][i]
    else:
        df.iloc[i, df.columns == 'average_price' ] = df['average_price'][i-1]
df.drop('prev_cum_qty', axis=1)

这种方法的一个优点是，如果有新的购买，它也会起作用 在 cum_qty 变为零之前。例如，假设有一个新的购买 5 的价格为 3，即在处理之前运行以下行 数据:

# Add more data, exemplifying a different situation
df = df.append({'quantity': 5, 'price': 3}, ignore_index=True)

我希望得到以下结果:

   quantity  price  amount  cum_qty  average_price
0         1      5       5        1            5.0
1        -1      6      -6        0            5.0
2         2      3       6        2            3.0
3        -1      2      -2        1            3.0
4        -1      4      -4        0            3.0
5         1      2       2        1            2.0
6         1      3       3        2            2.5
7         1      4       4        3            3.0
8        -2      5     -10        1            3.0
9         5      3      15        6            3.0 # Not 4.0

也就是说，由于还有 1 件商品以价格 3 购买，cum_qty 现在是 6，平均价格仍然是 3。