Блог

У меня есть набор данных, подобный:

     Id   Status

    1     0
    1     0
    1     0
    1     0
    1     1
    2     0
    1     0 # --> gets removed since this row appears after id 1 already had a status of 1
    2     0
    3     0
    3     0
 

Я хочу удалить все строки идентификатора после того, как его статус стал 1, т. Е. Мой новый набор данных будет:

     Id   Status

    1     0
    1     0
    1     0
    1     0
    1     1
    2     0
    2     0
    3     0
    3     0
 

Я хочу узнать, как эффективно реализовать это вычисление, поскольку у меня очень большой (более 200 ГБ) набор данных.

Решение, которое у меня есть в настоящее время, состоит в том, чтобы найти индекс первого 1 и разрезать каждую группу таким образом. В случаях, когда не существует 1, верните группу без изменений:

 def remove(series):
    indexless = series.reset_index(drop=True)
    ones = indexless[indexless['Status'] == 1]
    if len(ones) > 0:
        return indexless.iloc[:ones.index[0]   1]

    else:
        return indexless

df.groupby('Id').apply(remove).reset_index(drop=True)
 

Однако это выполняется очень медленно, есть ли способ исправить это или альтернативно ускорить вычисления?

Первая идея — создать совокупную сумму для групп с логической маской, но также необходимо shift , чтобы избежать потери в первую очередь 1 :

 #pandas 0.24 
s = (df['Status'] == 1).groupby(df['Id']).apply(lambda x: x.shift(fill_value=0).cumsum())
#pandas below
#s = (df['Status'] == 1).groupby(df['Id']).apply(lambda x: x.shift().fillna(0).cumsum())
df = df[s == 0]
print (df)
   Id  Status
0   1       0
1   1       0
2   1       0
3   1       0
4   1       1
5   2       0
7   2       0
8   3       0
9   3       0
 

Другим решением является использование пользовательской лямбда-функции с Series.idxmax :

 def f(x):
    if x['new'].any():
        return x.iloc[:x['new'].idxmax() 1, :]
    else:
        return x

df1 = (df.assign(new=(df['Status'] == 1))
        .groupby(df['Id'], group_keys=False)
        .apply(f).drop('new', axis=1))
print (df1)
    Id  Status
0    1       0
1    1       0
2    1       0
3    1       0
4    1       1
5    2       0
8    2       0
9    3       0
10   3       0
 

Или немного измененное первое решение — фильтруйте только группы с 1 помощью и применяйте решение только там:

 m = df['Status'].eq(1)
ids = df.loc[m, 'Id'].unique()
print (ids)
[1]

m1 = df['Id'].isin(m)
m2 = (m[m1].groupby(df['Id'])
            .apply(lambda x: x.shift(fill_value=0).cumsum())
            .eq(0))

df = df[m2.reindex(df.index, fill_value=True)]
print (df)
    Id  Status
0    1       0
1    1       0
2    1       0
3    1       0
4    1       1
5    2       0
8    2       0
9    3       0
10   3       0
 

Давайте начнем с этого набора данных.

 l =[[1,0],[1,0],[1,0],[1,0],[1,1],[2,0],[1,0], [2,0], [2,1],[3,0],[2,0], [3,0]]
df_ = pd.DataFrame(l, columns = ['id', 'status'])
 

Мы найдем индекс status = 1 для каждого идентификатора.

 status_1_indice = df_[df_['status']==1].reset_index()[['index', 'id']].set_index('id')

    index
id  
1   4
2   8
 

Теперь мы объединяемся с df_ status_1_indice

 join_table  = df_.join(status_1_indice, on='id').reset_index().fillna(np.inf)
 

Обратите .fillna(np.inf) внимание на идентификаторы, которые не имеют status=1. Результат:

     level_0 id  status  index
0   0   1   0   4.000000
1   1   1   0   4.000000
2   2   1   0   4.000000
3   3   1   0   4.000000
4   4   1   1   4.000000
5   5   2   0   8.000000
6   6   1   0   4.000000
7   7   2   0   8.000000
8   8   2   1   8.000000
9   9   3   0   inf
10  10  2   0   8.000000
11  11  3   0   inf
 

Требуемый фрейм данных может быть получен с помощью:

 join_table.query('level_0 <= index')[['id', 'status']]
 

Вместе:

 status_1_indice = df_[df_['status']==1].reset_index()[['index', 'id']].set_index('id')
join_table  = df_.join(status_1_indice, on='id').reset_index().fillna(np.inf)
required_df = join_table.query('level_0 <= index')[['id', 'status']]


   id   status
0   1   0
1   1   0
2   1   0
3   1   0
4   1   1
5   2   0
7   2   0
8   2   1
9   3   0
11  3   0
 

Я не могу ручаться за производительность, но это более прямолинейно, чем рассматриваемый метод.