优化针对大量输入值的Pandas DataFrame条件过滤方法

针对你提到的300万行DataFrame、数千个输入值的区间匹配需求，核心优化思路是利用X列已排序的特性，通过二分查找缩小候选范围，避免每次全表扫描，以下是具体实现方案：

优化原理 #

原循环方案每次对全表执行input >= X和input <= Y的判断，时间复杂度为O(m*n)（m为输入数量，n为DataFrame行数），300万行+5000个输入的总运算量达1.5e10次，效率极低。

优化后利用X列排序的特性：

1. 通过二分查找快速定位所有X <= input的行（O(log n)时间）
2. 仅在这些候选行中筛选Y >= input的结果（O(k)时间，k为候选行数，远小于n）
   总时间复杂度降至O(m*(log n + k))，效率提升几个数量级。

具体实现代码 #

1. 数据预处理（转numpy数组提升速度） #

先将DataFrame的X、Y列及索引转为numpy数组，比Pandas Series的向量运算更快：

import pandas as pd
import bisect
import numpy as np

# 模拟300万行测试数据（X已排序，Y > X）
np.random.seed(42)
X = np.sort(np.random.randint(0, 1000000, size=3000000))
Y = X + np.random.randint(1, 1000, size=3000000)
df = pd.DataFrame({"X": X, "Y": Y})

# 提取numpy数组，加速后续运算
X_np = df["X"].values
Y_np = df["Y"].values
indices_np = df.index.values

2. 单输入的高效匹配函数 #

def find_matching_indices(input_val):
    # 二分查找找到第一个X > input_val的位置，前面所有行满足X <= input_val
    right_idx = bisect.bisect_right(X_np, input_val)
    # 在候选行中筛选Y >= input_val的结果
    mask = Y_np[:right_idx] >= input_val
    # 返回对应的索引
    return indices_np[:right_idx][mask]

3. 批量处理数千个输入 #

ind_list = []
for input_val in range(1, 5000):
    inds = find_matching_indices(input_val)
    ind_list.append(inds)

4. 验证结果正确性（用示例数据测试） #

# 用你提供的示例数据验证
df_sample = pd.DataFrame({"X":[0,2,3,6,13], "Y":[10,12,16,8,22]})
X_np_sample = df_sample["X"].values
Y_np_sample = df_sample["Y"].values
indices_np_sample = df_sample.index.values

def find_matching_indices_sample(input_val):
    right_idx = bisect.bisect_right(X_np_sample, input_val)
    mask = Y_np_sample[:right_idx] >= input_val
    return indices_np_sample[:right_idx][mask]

print(find_matching_indices_sample(13))
# 输出：[2 4]，与原方案结果完全一致

额外优化建议 #

如果输入数量极大，还可以使用np.searchsorted批量计算所有输入的候选范围边界，进一步减少循环开销：

inputs = np.arange(1, 5000)
# 批量获取每个输入对应的候选行右边界
right_indices = np.searchsorted(X_np, inputs, side='right')

ind_list = []
for idx, input_val in zip(right_indices, inputs):
    mask = Y_np[:idx] >= input_val
    ind_list.append(indices_np[:idx][mask])

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