Для каждой строки в Dataframe проверьте все остальные строки на условие

Question

У меня есть таблица, подобная следующей:

df = pandas.DataFrame(
[[1,datetime.datetime(2018, 1, 1), datetime.datetime(2019, 1, 1),'Joe'], 
 [2,datetime.datetime(1999, 1, 1), datetime.datetime(2019, 1, 1),'James'],
 [3,datetime.datetime(1980, 1, 1), datetime.datetime(2000, 1, 1),'Jack'],
 [4,datetime.datetime(1967, 1, 1), datetime.datetime(1975, 1, 1),'Jim']],
columns=['PERSON ID', 'START DATE', 'END DATE','NAME']

Я хочу получить список коллег и подсчитать их для каждого человека на основе его начальной и конечной даты. Ожидаемый результат -

    Name Number_of_colleagues List_of_colleagues
    Joe 1 [James]
    James 2 [Joe, Jack]
    Jack 1 [James]
    Jim 0 []

Любая рекомендация, как это сделать? Я попытался использовать вложенный цикл для итерации по каждой строке, чтобы найти его. Это работает, но очень медленно для 20000 строк.

Answer 1

С пандами 0.25.3 вы можете использовать pd.IntervalIndex.overlaps :

df.index = pd.IntervalIndex.from_arrays(df['START DATE'], df['END DATE'])

df['Num of Colleagues'] = [sum(df.index.overlaps(i))-1 for i in df.index]
print(df)

Вывод:

                          PERSON ID START DATE   END DATE   NAME  Num of Colleagues
(2018-01-01, 2019-01-01]          1 2018-01-01 2019-01-01    Joe                  1
(1999-01-01, 2019-01-01]          2 1999-01-01 2019-01-01  James                  2
(1980-01-01, 2000-01-01]          3 1980-01-01 2000-01-01   Jack                  1
(1967-01-01, 1975-01-01]          4 1967-01-01 1975-01-01    Jim                  0

И мы можем добавить списокКоллеги, как это:

df['List of colleagues'] = [df.loc[df.index.overlaps(i), 'NAME'].tolist() for i in df.index]

Вывод:

                          PERSON ID START DATE   END DATE   NAME  Num of Colleagues  List of colleagues
(2018-01-01, 2019-01-01]          1 2018-01-01 2019-01-01    Joe                  1        [Joe, James]
(1999-01-01, 2019-01-01]          2 1999-01-01 2019-01-01  James                  2  [Joe, James, Jack]
(1980-01-01, 2000-01-01]          3 1980-01-01 2000-01-01   Jack                  1       [James, Jack]
(1967-01-01, 1975-01-01]          4 1967-01-01 1975-01-01    Jim                  0               [Jim]

И, исключены "Я" в этой строке:

df['List of colleagues'] = [df.loc[df.index.overlaps(i) & (df.index != i), 'NAME'].tolist()  for i in df.index]

Вывод:

                          PERSON ID START DATE   END DATE   NAME  Num of Colleagues List of colleagues
(2018-01-01, 2019-01-01]          1 2018-01-01 2019-01-01    Joe                  1            [James]
(1999-01-01, 2019-01-01]          2 1999-01-01 2019-01-01  James                  2        [Joe, Jack]
(1980-01-01, 2000-01-01]          3 1980-01-01 2000-01-01   Jack                  1            [James]
(1967-01-01, 1975-01-01]          4 1967-01-01 1975-01-01    Jim                  0                 []