Pandas DataFrame: столбец группировки по группам, сортировка по дате и по усечению по условию

Question

У меня есть Pandas DataFrame, который выглядит примерно так:

import pandas as pd

df = pd.DataFrame([['a', '2018-09-30 00:03:00', 'that is a glove'],
                   ['b', '2018-09-30 00:04:00', 'this is a glove'],
                   ['b', '2018-09-30 00:09:00', 'she has ball'],
                   ['a', '2018-09-30 00:05:00', 'they have a ball'],
                   ['a', '2018-09-30 00:01:00', 'she has a shoe'],
                   ['c', '2018-09-30 00:04:00', 'I have a baseball'],
                   ['a', '2018-09-30 00:02:00', 'this is a hat'],
                   ['a', '2018-09-30 00:06:00', 'he has no helmet'],
                   ['b', '2018-09-30 00:11:00', 'he has no shoe'],
                   ['c', '2018-09-30 00:02:00', 'we have a hat'],
                   ['a', '2018-09-30 00:04:00', 'we have a baseball'],
                   ['c', '2018-09-30 00:06:00', 'they have no glove'],
                   ], 
                  columns=['id', 'time', 'equipment'])


   id                 time           equipment
0   a  2018-09-30 00:03:00     that is a glove
1   b  2018-09-30 00:04:00     this is a glove
2   b  2018-09-30 00:09:00        she has ball
3   a  2018-09-30 00:05:00    they have a ball
4   a  2018-09-30 00:01:00      she has a shoe
5   c  2018-09-30 00:04:00   I have a baseball
6   a  2018-09-30 00:02:00       this is a hat
7   a  2018-09-30 00:06:00    he has no helmet
8   b  2018-09-30 00:11:00      he has no shoe
9   c  2018-09-30 00:02:00       we have a hat
10  a  2018-09-30 00:04:00  we have a baseball
11  c  2018-09-30 00:06:00  they have no glove

Что я хотел бы сделать, это groupby id и в каждой группе отсортировать по time изатем верните каждую строку вплоть до строки, содержащей слово «ball». До сих пор я могу группировать и сортировать:

df.groupby('id').apply(lambda x: x.sort_values(['time'], ascending=True)).reset_index(drop=True)


   id                 time           equipment
0   a  2018-09-30 00:01:00      she has a shoe
1   a  2018-09-30 00:02:00       this is a hat
2   a  2018-09-30 00:03:00     that is a glove
3   a  2018-09-30 00:04:00  we have a baseball
4   a  2018-09-30 00:05:00    they have a ball
5   a  2018-09-30 00:06:00    he has no helmet
6   b  2018-09-30 00:04:00     this is a glove
7   b  2018-09-30 00:09:00        she has ball
8   b  2018-09-30 00:11:00      he has no shoe
9   c  2018-09-30 00:02:00       we have a hat
10  c  2018-09-30 00:04:00   I have a baseball
11  c  2018-09-30 00:06:00  they have no glove

Однако я хочу, чтобы вывод выглядел следующим образом:

   id                 time           equipment
0   a  2018-09-30 00:01:00      she has a shoe
1   a  2018-09-30 00:02:00       this is a hat
2   a  2018-09-30 00:03:00     that is a glove
3   a  2018-09-30 00:04:00  we have a baseball
4   a  2018-09-30 00:05:00    they have a ball
6   b  2018-09-30 00:04:00     this is a glove
7   b  2018-09-30 00:09:00        she has ball

Обратите внимание, что в группе c нет строк, возвращаемых св нем нет строк со словом «мяч». Группа c имеет слово «бейсбол», но это не то соответствие, которое мы ищем. Точно так же обратите внимание, что группа a не останавливается в ряду «бейсбол», так как мы останавливаемся в ряду с «мячом». Каков наиболее эффективный способ сделать это как с точки зрения скорости, так и с точки зрения памяти?

Answer 1

Вот мой подход:

# as the final expected output is sorted by id and time
# we start by doing so to the whole data
df = df.sort_values(['id','time'])

# mark the rows containing the word `ball`
has_ball = (df.equipment.str.contains(r'\bball\b') )

# cumulative number of rows with `ball` in the group
s = has_ball.groupby(df['id']).cumsum()

# there must be row with `ball`
valid_groups = has_ball.groupby(df['id']).transform('max')

print(df[valid_groups &
         (s.eq(0) |              # not containing `ball` before the first
         (s.eq(1) & has_ball)    # first row containing `ball`
         )
        ]  
     )

Вывод:

   id                time           equipment
4   a 2018-09-30 00:01:00      she has a shoe
6   a 2018-09-30 00:02:00       this is a hat
0   a 2018-09-30 00:03:00     that is a glove
10  a 2018-09-30 00:04:00  we have a baseball
3   a 2018-09-30 00:05:00    they have a ball
1   b 2018-09-30 00:04:00     this is a glove
2   b 2018-09-30 00:09:00        she has ball

Answer 2

Продолжая то, что вы сделали:

new_df = df.groupby('id').apply(lambda x: x.sort_values(['time'], ascending=True)).reset_index(drop=True)

new_df["mask"] = new_df.groupby("id").apply(lambda x: x["equipment"].str.contains(r"\bball\b",regex=True)).reset_index(drop=True)

result = (new_df.groupby("id").apply(lambda x : x.iloc[:x.reset_index(drop=True)["mask"].
                                     idxmax()+1 if x["equipment"].str.contains(r"\bball\b",regex=True).any() else 0])
          .reset_index(drop=True).drop("mask",axis=1))

print (result)

#
  id                 time           equipment
0  a  2018-09-30 00:01:00      she has a shoe
1  a  2018-09-30 00:02:00       this is a hat
2  a  2018-09-30 00:03:00     that is a glove
3  a  2018-09-30 00:04:00  we have a baseball
4  a  2018-09-30 00:05:00    they have a ball
5  b  2018-09-30 00:04:00     this is a glove
6  b  2018-09-30 00:09:00        she has ball
7  d  2018-09-30 00:06:00       I have a ball