如何为Pandas DataFrame生成open_lvl_value列：统计前置相交行数总和

需求说明 #

现有包含open_time（开始时间）和close_time（结束时间）列的Pandas DataFrame，需新增open_lvl_value列，该列每行的值为当前行及之前所有行中，时间范围与当前行相交的总行数。

示例数据 #

原始数据代码 #

import pandas as pd

data = [[40, 110],[50, 110],[70, 110],[90, 110],
        [100, 110],[40, 70],[50, 70],[50, 70],
        [60, 70],[60, 70],[100, 110],[130, 180],
        [150, 180],[150, 180]]
df = pd.DataFrame(data, columns=['open_time','close_time'])

原始数据输出 #

open_time  close_time
0          40         110
1          50         110
2          70         110
3          90         110
4         100         110
5          40          70
6          50          70
7          50          70
8          60          70
9          60          70
10        100         110
11        130         180
12        150         180
13        150         180

预期结果 #

open_time  close_time  open_lvl_value
0          40         110               1
1          50         110               2
2          70         110               3
3          90         110               4
4         100         110               5
5          40          70               4
6          50          70               5
7          50          70               6
8          60          70               7
9          60          70               8
10        100         110               5
11        130         180               1
12        150         180               2
13        150         180               3

解决方案 #

基础实现（适用于小数据集） #

利用apply遍历每行，计算当前行及之前行中满足区间相交条件的行数：

# 计算open_lvl_value列
df['open_lvl_value'] = df.apply(
    lambda row: ((df.loc[:row.name, 'open_time'] <= row['close_time']) & 
                 (row['open_time'] <= df.loc[:row.name, 'close_time'])).sum(),
    axis=1
)

print(df)

代码解释 #

• 区间相交判断条件：对于两个闭区间[a1, b1]和[a2, b2]，当a1 <= b2且a2 <= b1时，两个区间相交。
• 针对每行，筛选出当前行及之前的所有行，逐一判断是否满足相交条件，最后对满足条件的行数求和，得到open_lvl_value的值。

优化实现（适用于大数据集） #

对于数据量较大的场景，使用Numpy广播实现向量化计算，比apply效率更高：

import numpy as np

open_times = df['open_time'].values
close_times = df['close_time'].values

# 创建相交判断矩阵，matrix[i,j]表示第j行与第i行的时间范围是否相交
intersection_matrix = (open_times[:, np.newaxis] <= close_times) & (open_times <= close_times[:, np.newaxis])

# 对每行i，求和前i+1列（即当前行及之前所有行的相交数量）
df['open_lvl_value'] = np.array([intersection_matrix[i, :i+1].sum() for i in range(len(df))])

print(df)

内容的提问来源于stack exchange，提问作者PiDesign Interior