基于将另一列中的值与阈值进行比较来为数据帧行分配值的有效方法？

Question

我有一个包含 140,000 行的大型数据框。其中一列中有一个值，我需要将其与两个阈值进行比较，然后分配一个索引。然后，我还在另一个索引中找到最接近的对应值，但该索引特定于先前的任务。

作为示例，这是一个类似的任务： 我有一个时间序列 ts，温度为“temp”，湿度为“h”。

1. 根据两个温度阈值，选择“类别”0、1 或 2。
2. 存在三个湿度数组，每一类一个。查找 ts 的特定行中与湿度值最接近的值的索引。笔记;在示例中，我没有将湿度数组 (h_ref) 传递给函数，它们是全局可用的。

我可以通过以下方式做到这一点：

def temp_and_h(ts):
    temp_i = ts['temp']
    h_i = ts['h']
    
    if temp_i > 2:
        class=0
    elif temp_i < -2:
        class=1
    else:
        class=2

    # find the index of closest value in the correct (class-specific) array in "h_ref"
    diffs = [abs(x - h_i) for x in h_ref[class]]
    h_idx = diffs.index(min(diffs))

    return (class,h_idx)

这只是通过以下方式调用：

ts[['temp_idx', 'h_idx']] = ts.apply(stability_and_TI, axis=1)

但是，在我的机器上，这大约需要 670 秒。我需要尽可能加快速度。我还可以尝试哪些其他方法？我尝试使用 np.vectorize，但内存占用太大（>100GB）。

Answer 1

您可以使用

pd.cut

和

groupby.apply

:

进行矢量化

ts['temp_idx'] = pd.cut(ts['temp'],
                        bins=[-np.inf, -2, 2, np.inf],
                        labels=[1, 2, 0])

ts['h_idx'] = (ts.groupby('temp_idx', group_keys=False)['h']
                 .apply(lambda g: pd.Series((g.to_numpy()[:,None]-h_ref[g.name]
                                             ).argmin(axis=1),
                                            index=g.index
                                            )
                       )
              )

虚拟数据示例：

                         temp         h temp_idx  h_idx
2024-01-01 00:00:00  0.507979  0.029876        2      0
2024-01-01 01:00:00  2.081478  0.456833        0      1
2024-01-01 02:00:00 -2.090953  0.649144        1      4
2024-01-01 03:00:00  0.108276  0.278487        2      0
2024-01-01 04:00:00  3.929470  0.676255        0      1
2024-01-01 05:00:00  3.962931  0.590863        0      1
2024-01-01 06:00:00 -3.744147  0.023982        1      4
2024-01-01 07:00:00 -2.927571  0.558854        1      4
2024-01-01 08:00:00 -4.485328  0.259252        1      4
2024-01-01 09:00:00 -0.591902  0.415101        2      0

可重复输入：

np.random.seed(3)
N = 10
ts = pd.DataFrame({'temp': np.random.uniform(-5, 5, N),
                   'h': np.random.uniform(0, 1, N),
                  }, index=pd.date_range('2024-01-01', freq='h', periods=N))

h_ref = np.random.random((5, 5))