Python pandas：对分组值 x 的前 n 行求和

Question

给定示例数据

data = {'id': ['1', '1', '1', '2', '2', '3', '3', '3', '3', '4', '1', '4', '3', '2', '4', '3', '2', '3', '3', '4'],\
            'price': [5,5,3,15,4,7,18,22,10,19,21,15,31,5,41,17,18,2,10,1],\
            'quantity': [4,1,1,2,1,4,6,55,1,22,34,5,11,22,3,1,6,1,7,9]}


df = pd.DataFrame(data)
print(df.head(10))

id  price   quantity
1   5       4
1   5       1
1   3       1
2   15      2
2   4       1
3   7       4
3   18      6
3   22      55
3   10      1
4   19      22

我想对按“id”列分组的前 n 行求和，按价格排序，其中 n 是该 id 总和数量的 15%。
我目前正在使用 iterrows 循环来获得想要的结果，但它非常慢。 如下所示，它不是很漂亮。

unique = df['id'].unique()
list1 = []
for u in unique:
    idf = df[df['id'] == u]
    idf.sort_values(by=['price'], inplace=True)
    summed_quantity = idf['quantity'].sum()
    percent = math.ceil(summed_quantity * 0.15)
    quant = 0
    sum = 0
    for i, row in idf.iterrows():
            if row['quantity'] + quant < percent:
                quant = quant + row['quantity']
                sum = sum + (row['quantity'] * row['price'])
            if row['quantity'] + quant >= percent:
                quant = quant + row['quantity']
                sum = sum + (row['quantity'] * row['price'])
                marketPrice = sum / quant
                list1.append(marketPrice)
                break

我怎样才能让这个 function 更有效率？
我一直试图让它与 groupby 一起工作，但到目前为止我还没有取得任何进展。

编辑：
总而言之，我对不同的解决方案进行了计时。

timer2 的解决方案似乎也存在一些不一致，当 N 行变大（10k+）时，marketPrice 不正确。

N rows: 250
Timer1:
id
4     6355.0
6      981.0
32     927.0
40     191.0
55    7115.0
Timer1 took 0.031532s

Timer2:
    market_price
4         6355.0
6          981.0
32         927.0
40         191.0
55        7115.0
Timer2 took 0.550794s

N rows: 250000
Timer1:
id
1     712.667707
2    1427.281561
3    1237.422608
4     536.492775
5     712.897859
Timer1 took 0.903622s

Timer2:
   market_price
1   1097.579822
2   1770.665625
3   2703.163577
4    688.552706
5    932.140983
Timer2 took 22.406997s

代码：

import timeit
import numpy as np
import pandas as pd


##Timer1
def cumsum(df):
    # sort the values by id and price first
    df.sort_values(['id', 'price'], ascending=True, inplace=True)
    df.reset_index(drop=True, inplace=True)

    # calculate the cumsum of quantity, and sale_amount, group by id
    df['quant_cum'] = df.groupby('id')['quantity'].cumsum()
    df['sale_amount'] = df['price'] * df['quantity']
    df['sale_amount_cum'] = df.groupby('id')['sale_amount'].cumsum() 

    # calculate the market_price with sale_amount_cum/quant_cum
    df['market_price'] = df['sale_amount_cum'] / df['quant_cum']

    # calculate threshold
    df['quantity_15'] = np.ceil(df.groupby('id')['quantity'].transform('sum') * 0.15)


    # finally, get the first row which is larger then the threshold, and output market_price
    cond = df.quant_cum >= df.quantity_15
    result = df[cond].groupby('id')['market_price'].apply(lambda x: x.iloc[0])
    print(result.head(5))
###



###Timer2    
def cumsumiloc(df):
    df['total_price'] = df['quantity'] * df['price']
    res = dict(process_df(df))
    result = pd.DataFrame.from_dict(res, orient='index', columns=['market_price'])
    print(result.head(5))


def process_df(df): 
    for _id, subdf in df.groupby('id'):
        # get cumulative sum
        csdf = subdf[['total_price', 'quantity']].sort_values(['total_price','quantity']).cumsum().reset_index(drop=True)

        # calculate cuttoff quantity
        cutoff_quantity = (0.15) * csdf['quantity'].iloc[-1] - 2e-10
        # get index where cuttoff quantity would belong to
        idx = csdf['quantity'].searchsorted(cutoff_quantity)        

        # calculate mean price by dividing (cumsummed) price by (cumsummed) quantity
        mean_price = csdf['total_price'].iloc[idx] / csdf['quantity'].iloc[idx]

        # yield id and mean price tuple for easy dictionary generation
        yield (_id, mean_price)
###


##Timer3
def function3(g):
    g = g.sort_values('price').reset_index(drop=True)

    # calculate quantity threshold
    threshold = math.ceil(g.quantity.sum() * 0.15)

    # find the quantity cutoff index
    idxmax = g.quantity.cumsum().ge(threshold).idxmax()

    # extract the rows below the cutoff
    lowest_15_percent = g.loc[:idxmax]

    # calculate the overall price
    total_price = lowest_15_percent.quantity.mul(lowest_15_percent.price).sum()

    # calculate the average price
    return total_price / lowest_15_percent.quantity.sum()  

def marketPrice(df):
    resdf = df.groupby('id').apply(function3)
    print(resdf.head(5))
###

##Timer4
def function4(df):
    unique = df['id'].unique()
    list1 = []
    for u in unique:
        idf = df[df['id'] == u]
        idf = idf.sort_values(by=['price'])
        summed_quantity = idf['quantity'].sum()
        percent = math.ceil(summed_quantity * 0.15)
        quant = 0
        sum = 0
        for i, row in idf.iterrows():
                if row['quantity'] + quant < percent:
                    quant = quant + row['quantity']
                    sum = sum + (row['quantity'] * row['price'])
                elif row['quantity'] + quant >= percent:
                    quant = quant + row['quantity']
                    sum = sum + (row['quantity'] * row['price'])
                    marketPrice = sum / quant
                    list1.append({'id': u, 'marketPrice': marketPrice})
                    break
    listdf = pd.DataFrame(list1)
    print(listdf.head(5))

for i in [250, 2500, 250000, 2500000]:   
    dfrnd = pd.DataFrame(np.random.randint(1,10000,size=(i, 3)), columns=['id','price','quantity'])
    print("N rows: " + str(len(dfrnd)))
    t = timeit.Timer("cumsum(dfrnd)", globals=globals())
    time = t.timeit(1)
    print("Timer1 took %fs\n" % (time,))

    t2 = timeit.Timer("cumsumiloc(dfrnd)", globals=globals())
    time2 = t2.timeit(1)
    print("Timer2 took %fs\n" % (time2,))

    t3 = timeit.Timer("marketPrice(dfrnd)", globals=globals())
    time3 = t3.timeit(1)
    print("Timer3 took %fs\n" % (time3,))

    t4 = timeit.Timer("function4(dfrnd)", globals=globals())
    time4 = t4.timeit(1)
    print("Timer4 took %fs\n" % (time4,))

Answer 1

您可以使用groupby尝试以下操作以达到相同的结果：

import math

def lowest_15_percent_avg_price(g):
  g = g.sort_values('price').reset_index(drop=True)

  # calculate quantity threshold
  threshold = math.ceil(g.quantity.sum() * 0.15)

  # find the quantity cutoff index
  idxmax = g.quantity.cumsum().ge(threshold).idxmax()

  # extract the rows below the cutoff
  lowest_15_percent = g.loc[:idxmax]

  # calculate the overall price
  total_price = lowest_15_percent.quantity.mul(lowest_15_percent.price).sum()

  # calculate the average price
  return  total_price / lowest_15_percent.quantity.sum()

print(df.groupby('id').apply(lowest_15_percent_avg_price))
#id
#1    4.666667
#2    4.956522
#3    8.461538
#4    1.000000
#dtype: float64

请参阅此链接以与 for 循环结果进行比较。 请注意，在您的原始 for 循环中，第二个 if 语句需要是elif ，否则可能会重复计算截止条目。

Answer 2

我只是尝试自己解决您的问题，这是我编写的代码。

我认为在性能方面的两个主要收益首先是使用.cumsum()因为这是在 C 中实现的，并且比手动添加数字要快得多。 其次（与其他答案相比）使用.searchsorted()因为据我所知，这比任何替代方法都快。

data = {'id': ['1', '1', '1', '2', '2', '3', '3', '3', '3', '4', '1', '4', '3', '2', '4', '3', '2', '3', '3', '4'],\
            'price': [5,5,3,15,4,7,18,22,10,19,21,15,31,5,41,17,18,2,10,1],\
            'quantity': [4,1,1,2,1,4,6,55,1,22,34,5,11,22,3,1,6,1,7,9]}
df = pd.DataFrame(data)

df['total_price'] = df['quantity'] * df['price']


def process_df(df): 
    for _id, subdf in df.groupby('id'):
        # get cumulative sum
        csdf = subdf[['total_price', 'quantity']].sort_values(['total_price','quantity']).cumsum().reset_index(drop=True)
        
        # calculate cuttoff quantity
        cutoff_quantity = (0.15) * csdf['quantity'].iloc[-1] - 2e-10
        # get index where cuttoff quantity would belong to
        idx = csdf['quantity'].searchsorted(cutoff_quantity)        
        
        # calculate mean price by dividing (cumsummed) price by (cumsummed) quantity
        mean_price = csdf['total_price'].iloc[idx] / csdf['quantity'].iloc[idx]
        
        # yield id and mean price tuple for easy dictionary generation
        yield (_id, mean_price)

dict(process_df(df))

为了实现完全相同的截止行为，我从cutoff_quantity中减去一个非常小的数字（ 2e-10 ）。 结果list1如下：

{'1': 4.666666666666667,
 '2': 15.777777777777779,
 '3': 9.071428571428571,
 '4': 1.0}

我可能误解了你的截止行为，为了检查你自己，你可以在yield语句之前粘贴这个块来检查哪个决定是以什么方式做出的：

        ### start of debug block
        print('Groupby-SubDF for id', _id)
        display(subdf) # use print(subdf) if you don't use jupyternotebook
        print('CumsumDF for id', _id)
        display(csdf) # use print(subdf) if you don't use jupyternotebook
        print('cutoff quantity:', cutoff_quantity, end = ', ')
        print("cutoff index:", idx, end = ', ')
        print('mean_price = ', csdf['total_price'].iloc[idx], '/', csdf['quantity'].iloc[idx], '=', mean_price, end='\n\n')
        ### end of debug block

这是前 2 个 ID 的决策。

Answer 3

df = pd.DataFrame(data)

# sort the values by id and price first
df.sort_values(['id', 'price'], ascending=True, inplace=True)
df.reset_index(drop=True, inplace=True)

# calculate the cumsum of quantity, and sale_amount, group by id
df['quant_cum'] = df.groupby('id')['quantity'].cumsum()
df['sale_amount'] = df['price'] * df['quantity']
df['sale_amount_cum'] = df.groupby('id')['sale_amount'].cumsum() 

# calculate the market_price with sale_amount_cum/quant_cum
df['market_price'] = df['sale_amount_cum'] / df['quant_cum']

# calculate threshold
df['quantity_15'] = np.ceil(df.groupby('id')['quantity'].transform('sum') * 0.15)


# finally, get the first row which is larger then the threshold, and output market_price
cond = df.quant_cum >= df.quantity_15
result = df[cond].groupby('id')['market_price'].apply(lambda x: x.iloc[0])


# print(df.head())
#   id  price  quantity  quant_cum  sale_amount  sale_amount_cum  market_price  \
# 0  1      3         1          1            3                3      3.000000   
# 1  1      5         4          5           20               23      4.600000   
# 2  1      5         1          6            5               28      4.666667   
# 3  1     21        34         40          714              742     18.550000   
# 4  2      4         1          1            4                4      4.000000   

#    quantity_15  
# 0          6.0  
# 1          6.0  
# 2          6.0  
# 3          6.0  
# 4          5.0 

print(result)

# id
# 1    4.666667
# 2    4.956522
# 3    8.461538
# 4    1.000000
# Name: market_price, dtype: float64