Частный случай итерации через pandas фрейм данных

Question

Я хотел бы заполнить столбцы данных с разницей во времени между текущей и ближайшей временными метками типа «тип А» или «не тип А» соответственно, т. Е. Type_A = 1 или type_A = 0. Далее показано небольшое пример:

import numpy as np
import pandas as pd
from datetime import datetime

df = pd.DataFrame({'id':[1,2,3,4], 
                   'tmstmp':[datetime(2018,5,4,13,27,10), datetime(2018,5,3,13,27,10),
                             datetime(2018,5,2,13,27,10), datetime(2018,5,1,13,27,10)], 
                   'type_A':[0, 1, 0, 1],
                   'dt_A': [np.nan]*4,
                   'dt_notA': [np.nan]*4
                  })

(строки A и не A не обязательно чередуются, но столбец отметки времени уже отсортирован в порядке убывания). Я вычисляю разницу во времени между отметкой времени в текущей строке и следующей строке с type_A = 1 или type_A = 0, соответственно, путем итерации по целочисленному индексу строки и элементам доступа по этому целочисленному индексу и имени столбца:

keys = {1: 'dt_A', 0: 'dt_notA'}
ridx = 0
while ridx + 1 < df.shape[0]:
    ts1 = df.iloc[ridx]['tmstmp']
    ts2 = df.iloc[ridx + 1]['tmstmp']
    found = 0 if df.iloc[ridx + 1]['type_A'] == 0 else 1
    key = keys[found]
    df.loc[ridx, key] = (ts1 - ts2).total_seconds()/3600
    complement = 1 - found
    j = 2
    while ridx + j < df.shape[0] and df.iloc[ridx + j]['type_A'] != complement:
        j += 1
    if ridx + j < df.shape[0]:
        ts1 = df.iloc[ridx]['tmstmp']
        ts2 = df.iloc[ridx + j]['tmstmp']
        val = (ts1 - ts2).total_seconds()/3600
    else:
        val = np.nan
    df.loc[ridx, keys[complement]] = val
    ridx += 1

Итерации по фрейму данных "не рекомендуется" из соображений эффективности (см. Как перебирать строки в фрейме данных в Pandas? ), а использование целочисленных индексов еще меньше "pythoni c ", поэтому мой вопрос таков: в этом конкретном случае есть ли" лучший "(более эффективный, более питонийный c) способ перебирать кадр данных для достижения поставленной задачи? Большое спасибо за любые предложения или мысли!

Редактировать : входные и выходные кадры данных для небольшого примера - столбец dt_A содержит дельты времени между текущей строкой и следующей, которая имеет type_A = 1, dt_notA содержит дельты времени с ближайшей строкой, которая имеет type_A = 0.

input: 
   id              tmstmp  type_A  dt_A  dt_notA
0   1 2018-05-04 13:27:10       0   NaN      NaN
1   2 2018-05-03 13:27:10       1   NaN      NaN
2   3 2018-05-02 13:27:10       0   NaN      NaN
3   4 2018-05-01 13:27:10       1   NaN      NaN

вывод:

   id              tmstmp  type_A  dt_A  dt_notA
0   1 2018-05-04 13:27:10       0  24.0     48.0
1   2 2018-05-03 13:27:10       1  48.0     24.0
2   3 2018-05-02 13:27:10       0  24.0      NaN
3   4 2018-05-01 13:27:10       1   NaN      NaN

Answer 1

def next_value_index(l, i, val):
    """Return index of l where val occurs next from position i."""
    try:
        return l[(i+1):].index(val) + (i + 1)
    except ValueError:
        return np.nan

def next_value_indexes(l, val):
    """Return for each position in l next-occurrence-indexes of val in l."""
    return np.array([next_value_index(l, i, val) for i, _ in enumerate(l)])

def nan_allowing_access(df, col, indexes):
    """Return df[col] indexed by indexes. A np.nan would cause errors.
    This function returns np.nan where index is np.nan."""
    idxs = np.array([idx if not np.isnan(idx) else 0 for idx in indexes])
    res = df[col].iloc[idxs]
    res[np.isnan(indexes)] = np.nan
    return res # NaT for timestamps

def diff_timestamps(dfcol1, dfcol2): # timestamp columns of pandas subtraction
    return [x - y for x, y in zip(list(dfcol1), list(dfcol2))]
    # this is not optimal in speed, but numpy did unwanted type conversions
    # problem is: np.array(df[tmstmp_col]) converts to `dtype='datetime64[ns]'`

def td2hours(timedelta): # convert timedelta to hours
    return timedelta.total_seconds() / 3600

def time_diff_to_next_val(df, tmstmp_col, col, val, converter_func, flip_subtraction=False):
    """
    Return time differences (timestamps are given in tmstmp_col column
    of the pandas data frame `df`) from the row's timestamp to the next
    time stamp of the row, which has in column `col` the next occurrence
    of value given in `val` in the data frame.

    converter_func is the function used to convert the timedelta
             value.
    flip_subtraction determines the order of subtraction: whether take current row's
             timestamp first or not when subtracting
    """
    next_val_indexes = next_value_indexes(df[col].tolist(), val)
    next_val_timestamps = nan_allowing_access(df, tmstmp_col, next_val_indexes)
    return [converter_func(x) for x in diff_timestamps(*(df[tmstmp_col], next_val_timestamps)[::(1-2*flip_subtraction)])]
    # `*(df[tmstmp_col], next_val_timestamps)[::(1-2*flip_subtraction)]`
    # flips the order of arguments when `flip_subtraction = True`

Применение функций:

df['dt_A'] = time_diff_to_next_val(df,'tmstmp', 'type_A', 1, converter_func = td2hours)
df['dt_notA'] = time_diff_to_next_val(df,'tmstmp', 'type_A', 0, converter_func = td2hours)

Тогда df становится:

   id              tmstmp  type_A  dt_A  dt_notA
0   1 2018-05-04 13:27:10       0  24.0     48.0
1   2 2018-05-03 13:27:10       1  48.0     24.0
2   3 2018-05-02 13:27:10       0  24.0      NaN
3   4 2018-05-01 13:27:10       1   NaN      NaN