Сопряжение со строками «вмещающие даты» в pandas

Question

У меня есть DataFrame с двумя типами задач, скажем Major и Minor, которые выполняются некоторыми агентами между датами:

ID  | StartDate           | EndDate             | Agent    | Task
-----------------------------------------------------------------
105 | 2020-01-01 12:00:00 | 2020-01-01 12:12:35 | Marty    | Major
106 | 2020-01-01 12:04:21 | 2020-01-01 12:09:08 | Wendy    | Major
109 | 2020-01-01 12:07:03 | 2020-01-01 12:11:48 | Marty    | Minor
114 | 2020-01-01 12:14:35 | 2020-01-01 12:19:10 | Wendy    | Major
117 | 2020-01-01 12:16:22 | 2020-01-01 12:16:41 | Wendy    | Minor
120 | 2020-01-01 12:17:03 | 2020-01-01 12:18:32 | Wendy    | Minor

Я хочу знать, когда агент выполнял второстепенную задачу во время также выполняя мажорное задание (хорошо делать только мажорное задание без Младшего параллельно). В приведенном ниже примере Незначительная задача с идентификатором 109 начинается после основной задачи 105 и заканчивается раньше, поэтому она считается. Если во время одной основной задачи выполняется несколько второстепенных задач, я рад просто оставить первую. В идеале, я хотел бы, чтобы дополнительный столбец сообщал мне, что задачи объединены в пару, например:

ID  | StartDate           | EndDate             | Agent    | Task  | PairedWith
-------------------------------------------------------------------------------
105 | 2020-01-01 12:00:00 | 2020-01-01 12:12:35 | Marty    | Major | 109
106 | 2020-01-01 12:04:21 | 2020-01-01 12:09:08 | Wendy    | Major |
109 | 2020-01-01 12:07:03 | 2020-01-01 12:11:48 | Marty    | Minor | 105
114 | 2020-01-01 12:14:35 | 2020-01-01 12:19:10 | Wendy    | Major | 117
117 | 2020-01-01 12:16:22 | 2020-01-01 12:16:41 | Wendy    | Minor | 114
120 | 2020-01-01 12:17:03 | 2020-01-01 12:18:32 | Wendy    | Minor | 114

Я ищу способ, не слишком медленный для этого. В данный момент я перебираю все строки с itertupes, а затем проверяю условия «даты включения и того же агента», что является самой медленной частью моего анализа. Был бы более быстрый путь, может быть, со стратегией apply ...?

Answer 1

Используйте groupby и shift.

from io import StringIO
import pandas as pd

data = StringIO("""ID|StartDate|EndDate|Agent|Task
105|2020-01-01 12:00:00|2020-01-01 12:12:35|Marty|Major
106|2020-01-01 12:04:21|2020-01-01 12:09:08|Wendy|Major
109|2020-01-01 12:07:03|2020-01-01 12:11:48|Marty|Minor
114|2020-01-01 12:14:35|2020-01-01 12:19:10|Wendy|Major
117|2020-01-01 12:16:22|2020-01-01 12:16:41|Wendy|Minor
120|2020-01-01 12:17:03|2020-01-01 12:18:32|Wendy|Minor""")

# read data and convert to datetime to allow comparison
df = pd.read_csv(data, sep='|')
df['StartDate'] = pd.to_datetime(df['StartDate'])
df['EndDate'] = pd.to_datetime(df['EndDate'])

# add columns for next start and id
df['NextStart'] = df.groupby('Agent')['StartDate'].shift(-1)
df['NextID'] = df.groupby('Agent')['ID'].shift(-1)

# if DropMinor is true, then next task started within window of this task
df['DropMinor'] = df['NextStart'] <= df['EndDate']

# find NextID where DropMinor is True and converts to int
exclude_ids = [int(i) for i in df.loc[df['DropMinor'], 'NextID']]

# exclude minor tasks within window
df.loc[~df.ID.isin(exclude_ids)]