Pandas новый фрейм данных на основе комбинаций всех значений в столбце

Question

Я собирал данные о местоположении автобусов в течение некоторого времени и хочу построить модель, предсказывающую, когда автобус прибудет на определенную остановку.

В самой простой форме у меня есть DataFrame, подобный этому:

import pandas as pd

df = pd.DataFrame({'station': ['Station 1', 'Station 2', 'Station 3', 'Station 4'], 
                    'arrival_time': ['10:00', '10:02', '10:03', '10:05']})
print(df)

     station arrival_time
0  Station 1        10:00
1  Station 2        10:02
2  Station 3        10:03
3  Station 4        10:05

Я хотел бы сопоставить время прибытия на каждую станцию ​​со временем прибытия на станцию ​​позже в поездке. Ожидаемый результат выглядит примерно так:

  station_prev arrival_time_prev station_next arrival_time_next
0    Station 1             10:00    Station 2             10:02
1    Station 2             10:02    Station 3             10:03
2    Station 3             10:03    Station 4             10:05
3    Station 1             10:00    Station 3             10:03
4    Station 2             10:02    Station 4             10:05
5    Station 1             10:00    Station 4             10:05

Я экспериментировал с df.shift (), и для одиночных DataFrames работает следующее.

import pandas as pd
import numpy as np

def combos(df):
    
    columns_prev = np.array(df.columns) + '_prev'
    columns_next = np.array(df.columns) + '_next'
 
    df_combo = pd.DataFrame()
    
    for i in range(1, df.shape[0]):
        df_prev = df.shift(i)
        df_prev.columns = columns_prev
        df_next = df.copy()
        df_next.columns = columns_next
        combo = pd.concat([df_prev, df_next], axis=1).dropna()
        df_combo = df_combo.append(combo, ignore_index=True)
    
    return df_combo

Однако это довольно медленно для DataFrames большего размера и регулярно ломается, когда я пытаюсь превратить его в более крупную функцию, которая собирает данные из многих поездок (я часто получаю ключевые ошибки, но не понимаю почему). Есть идеи, как сделать это более элегантно, эффективно и надежно? Заранее большое спасибо!

Answer 1

Преобразование «станции» в упорядоченный категориальный столбец:

df['station'] = pd.Categorical(df['station'], ordered=True).codes  

Теперь вы можете выполнить перекрестное соединение и фильтровать:

tmp = df.assign(key=1)
(tmp.merge(tmp, on='key', suffixes=('_prev', '_next'))
    .drop('key', 1)
    .query('station_prev < station_next'))

    station_prev arrival_time_prev  station_next arrival_time_next
1              0             10:00             1             10:02
2              0             10:00             2             10:03
3              0             10:00             3             10:05
6              1             10:02             2             10:03
7              1             10:02             3             10:05
11             2             10:03             3             10:05