通过 pandas 数据框进行迭代的特殊情况

Question

我想用当前时间戳之间的时间差填充数据框列 和“type A”或“not type A”最接近的时间戳，即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

出于效率原因，“不鼓励”对数据帧进行迭代（请参阅How to iterate over rows in a DataFrame in Pandas?） 并且使用整数索引甚至更少“pythonic”，所以我的问题是：在这种特殊情况下，是否有“更好”（更高效，更pythonic） 遍历数据框以实现给定任务的方法？ 非常感谢您的任何建议或想法！

编辑：小示例的输入和输出数据帧 - 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