如果前面的行满足条件，则 pandas dataframe 的最后 N 行的平均值

Question

我有一个 pandas dataframe 之类的

   index    start end  label
   0          2     5    0
   1          3     8    1
   2          4     8    0
   3          5     9    1
   4          6    10    0
   5          7    10    1
   6          8    11    1
   7          9    12    0

我想要一个新的专栏“意思”； 其中值是前几行label的平均值，条件为df['start']<df['end']

例子，

对于索引 1， df['mean'] = (df[0]['label']+ df[1]['label'])/2

对于索引 3， df['mean'] = (df[1]['label']+ df[2]['label']+ df[3]['label'])/3 ； 这里我们忽略索引 0，因为df[3]['start']<df[0]['end']条件不满足。

同样，对于索引 7， df['mean'] = (df[4]['label']+ df[5]['label']+ df[6]['label']+ df[7]['label'])/4 ; 至于索引 0,1,2,3； df[7]['start']<df[i]['end']条件不满足。

所以最终的 output 将是

  index    start end  label     mean
       0          2     5    0  0
       1          3     8    1  1/2
       2          4     8    0  1/3
       3          5     9    1  2/3
       4          6    10    0  2/4
       5          7    10    1  3/5
       6          8    11    1  3/4
       7          9    12    0  2/4

我正在尝试使用cumsum ； 但我不确定如何提出条件。

Answer 1

这是你的结果。

import numpy as np
mask_matrix = (
    (df.start.to_numpy().reshape(1,-1).T < df.end.to_numpy()) 
    & (df.index.to_numpy() <= np.arange(0,len(df)).reshape(1, -1).T)
)
df_add = pd.DataFrame(
    (np.matmul(
        (
            (mask_matrix) 
            
        ), df.label.to_numpy()
    )
) / (mask_matrix.sum(axis=-1)),
    columns = ["mean"]
)
df = pd.concat([df, df_add], axis=1)

当我们创建矩阵时，我们使用 O(n^2) 的额外空间。 希望这不是问题。 否则在使用矢量化计算时需要使用我个人不喜欢的循环。

一些额外的评论： df.start.to_numpy().reshape(1,-1).T < df.end.to_numpy() 基本上比较每一行的 start 低于 end 的位置。 这是结果：

array([[ True,  True,  True,  True,  True,  True,  True,  True,  True],
   [ True,  True,  True,  True,  True,  True,  True,  True,  True],
   [ True,  True,  True,  True,  True,  True,  True,  True,  True],
   [False,  True,  True,  True,  True,  True,  True,  True,  True],
   [False,  True,  True,  True,  True,  True,  True,  True,  True],
   [False,  True,  True,  True,  True,  True,  True,  True,  True],
   [False, False, False,  True,  True,  True,  True,  True, False],
   [False, False, False, False,  True,  True,  True,  True, False],
   [False, False, False, False,  True,  True,  True,  True, False]])

(df.index.to_numpy() <= np.arange(0,len_).reshape(1, -1).T) 将先前的结果限制为仅在当前结果之前的行。 这个面具看起来像这样：

array([[ True, False, False, False, False, False, False, False, False],
   [ True,  True, False, False, False, False, False, False, False],
   [ True,  True,  True, False, False, False, False, False, False],
   [ True,  True,  True,  True, False, False, False, False, False],
   [ True,  True,  True,  True,  True, False, False, False, False],
   [ True,  True,  True,  True,  True,  True, False, False, False],
   [ True,  True,  True,  True,  True,  True,  True, False, False],
   [ True,  True,  True,  True,  True,  True,  True,  True, False],
   [ True,  True,  True,  True,  True,  True,  True,  True,  True]])

最终 mask_matrix（前两个矩阵的元素乘法）看起来像这样

array([[ True, False, False, False, False, False, False, False, False],
       [ True,  True, False, False, False, False, False, False, False],
       [ True,  True,  True, False, False, False, False, False, False],
       [False,  True,  True,  True, False, False, False, False, False],
       [False,  True,  True,  True,  True, False, False, False, False],
       [False,  True,  True,  True,  True,  True, False, False, False],
       [False, False, False,  True,  True,  True,  True, False, False],
       [False, False, False, False,  True,  True,  True,  True, False],
       [False, False, False, False,  True,  True,  True,  True, False]])

现在将此 mask_matrix 乘以向量 df.label 几乎可以满足我们的需要。 只需要按元素除以 mask_matrix 中 True 的总和

Answer 2

这是一个性能较低的解决方案（在 Pandas 中通常应避免在每一行上循环），但希望可以将其作为一个起点，然后您可以对其进行优化：

df = pd.DataFrame([
    [2,5,0],
    [3,8,1],
    [4,8,0],
    [5,9,1],
    [6,10,0],
    [7,10,1],
    [8,11,1],
    [9,12,0]],columns=['start','end','label'])


for index, row in df.iterrows():

    if index == 0:
        df.at[index, 'cumulative_mean'] = 0

    else:
        current_row_start = row['start']
        previous_rows_as_df = df.loc[0:index] # create a DF which is all the previous rows
        
        for p_index, p_row in previous_rows_as_df.iterrows():
            if current_row_start < p_row['end']:
                previous_rows_as_df.at[p_index, 'include'] = True

        df.at[index, 'cumulative_mean'] = previous_rows_as_df[previous_rows_as_df['include'] == True]['label'].mean()

     

Answer 3

公平地说，我决定尝试比较三种方法：

1. 使用循环
2. 使用并行循环（numba）
3. 使用矩阵

这是代码

import pandas as pd
from numba import njit, prange

import numpy as np

from timeit import timeit

from pandas.testing import assert_frame_equal

big_df = pd.DataFrame(np.random.randint(0,100,size=(1000, 3)), columns=["start", "end", "label"])

def cond_cumsum_matrix(df):
    mask_matrix = (
        (df.start.to_numpy().reshape(1,-1).T < df.end.to_numpy()) 
        & (df.index.to_numpy() <= np.arange(0,len(df)).reshape(1, -1).T)
    )
    with np.errstate(divide='ignore', invalid='ignore'):
        df_add = pd.DataFrame(
            (np.matmul(
                (
                    (mask_matrix) 

                ), df.label.to_numpy()
            )
        ) / (mask_matrix.sum(axis=-1)),
            columns = ["mean"]
        )
    return df_add

def cond_cumsum_parallel_loop(df):
    @njit
    def numba_cond_cumsum_parallel_loop(label, start, end):
        cumsum = []        
        
        for i in prange(len(label)):
            running = 0
            count = 0
            for j in prange(i+1):
                if start[i] < end[j] :
                    running += label[j]
                    count += 1
            if count == 0:
                cumsum.append(np.nan)            
            else:
                cumsum.append(running/count)
        return cumsum
        
    return pd.DataFrame(
        numba_cond_cumsum_parallel_loop(
        df.label.to_numpy(),
        df.start.to_numpy(),
        df.end.to_numpy(),
    ), columns=["mean"],) 

def cond_cumsum_loop(df):
    start = df.start.tolist()
    end = df.end.tolist()
    label = df.label.tolist()
    
    cumsum = []                
    for index, row in df.iterrows():

        running = 0
        count = 0
        for j in range(index+1):
            if row.start < end[j] :
                running += label[j]
                count += 1
            
        if count == 0:
            cumsum.append(np.nan)
            
        else:
            cumsum.append(running/count)
        
    return pd.DataFrame(
        cumsum,
    columns=["mean"],)

assert_frame_equal(cond_cumsum_matrix(big_df), cond_cumsum_loop(big_df))
assert_frame_equal(cond_cumsum_matrix(big_df), cond_cumsum_parallel_loop(big_df))

repetitions = 5
print(f"cond_cumsum_loop runs {timeit(lambda: cond_cumsum_loop(big_df), number=repetitions)/repetitions} seconds")
print(f"cond_cumsum_parallel_loop runs {timeit(lambda: cond_cumsum_parallel_loop(big_df), number=repetitions)/repetitions} seconds")
print(f"cond_cumsum_matrix runs {timeit(lambda: cond_cumsum_matrix(big_df), number=repetitions)/repetitions} seconds")   

这是它给出的结果：

cond_cumsum_loop runs 1.2179410583339632 seconds
cond_cumsum_parallel_loop runs 0.07655967501923441 seconds
cond_cumsum_matrix runs 0.004219983238726854 seconds

当然，代码可以改进，因此比较并不理想，但无论如何，结论是尽管矩阵仍然在使用 O(n^2) 额外 memory 的情况下在性能上获胜，但并行循环仅使用 O(n ) 附加 memory。