python - 在 Pandas 中的特定值之前获取 n 行

Question

说，我有以下数据框:

import pandas as pd
dict = {'val':[3.2, 2.4, -2.3, -4.9, 3.2, 2.4, -2.3, -4.9, 2.4, -2.3, -4.9], 
        'label': [0, 2, 1, -1, 1, 2, -1, -1,1, 1, -1]} 
df = pd.DataFrame(dict) 
df
     val    label
0    3.2     0
1    2.4     2
2   -2.3     1
3   -4.9    -1
4    3.2     1
5    2.4     2
6   -2.3    -1
7   -4.9    -1
8    2.4     1
9   -2.3     1
10  -4.9    -1

我想在列标签中的 -1 值之前取每 n(例如 2)行。在给定的 df 中，首先 -1 出现在索引 3 处，我们取它之前的 2 行并删除索引 3，然后下一个 -1 出现在索引 6 处，我们再次保留之前的 2 行等等。所需的输出如下:

    val     label
1    2.4     2
2   -2.3     1
4    3.2     1
5    2.4     2
6   -2.3    -1
8    2.4     1
9   -2.3     1

感谢您的任何想法!

Answer 1

您可以获得 index values 然后获取前两行索引值:

idx = df[df.label == -1].index
filtered_idx = (idx-1).union(idx-2)
filtered_idx = filtered_idx[filtered_idx > 0]

df_new = df.iloc[filtered_idx]

输出:

   val  label
1  2.4      2
2 -2.3      1
4  3.2      1
5  2.4      2
6 -2.3     -1
8  2.4      1
9 -2.3      1

速度对比 for loop解决方案:

# create large df:
import numpy as np
df = pd.DataFrame(np.random.random((20000000,2)), columns=["val","label"])
df.loc[df.sample(frac=0.01).index, "label"] = - 1

def vectorized_filter(df):
    idx = df[df.label == -1].index
    filtered_idx = (idx -1).union(idx-2)
    df_new = df.iloc[filtered_idx]
    return df_new

def loop_filter(df):
    filter = df.loc[df['label'] == -1].index
    req_idx = []
    for idx in filter:
        if idx == 0:
            continue
        elif idx == 1:
            req_idx.append(idx-1)
        else:
            req_idx.append(idx-2)
            req_idx.append(idx-1)    
    req_idx = list(set(req_idx))
    df2 = df.loc[df.index.isin(req_idx)]
    return df2

%timeit vectorized_filter(df)
%timeit loop_filter(df)

vectorized runs ~20x faster on my machine