如何向量化优化基于两个NumPy数组的Pandas遍历循环？

你的循环代码每次都要切片DataFrame并过滤，当t_indices规模较大时，重复的切片和过滤会导致性能急剧下降。下面提供两种向量化优化方案，性能远优于循环实现：

方案一：用merge_asof实现高效顺序匹配 #

merge_asof是Pandas专为“按顺序找后续首个符合条件的行”设计的工具，时间复杂度为O(N log N)，适合大规模数据：

1. 先确保原DataFrame索引单调递增（若无序先排序），并把t_indices和t_lows整理成临时DataFrame：

import pandas as pd
import numpy as np

# 确保df索引单调递增
df = df.sort_index()

# 创建临时DataFrame，保留原始顺序避免匹配后打乱
temp_df = pd.DataFrame({
    't_index': t_indices,
    'low': t_lows,
    'original_order': np.arange(len(t_indices))
})

2. 执行merge_asof匹配，指定找t_index之后的首个符合条件的行：

# 临时DataFrame需按匹配列排序
temp_df = temp_df.sort_values('t_index')
# 把df的索引转为列，方便后续提取
df['orig_index'] = df.index

# 执行匹配：forward表示找后续行，allow_exact_matches=False排除t_index本身
result = pd.merge_asof(
    temp_df,
    df[['orig_index', 'Low']],
    left_on='t_index',
    right_on='orig_index',
    direction='forward',
    allow_exact_matches=False,
    condition=df['Low'] <= temp_df['low']
)

# 恢复原始顺序并提取结果索引
result = result.sort_values('original_order')
t_ixs = result['orig_index'].values

方案二：广播+argmax实现向量化查找 #

适合数据规模中等的场景，通过广播生成条件矩阵，再找首个符合条件的位置：

1. 先将df索引转为位置映射（若索引无序先排序）：

df = df.sort_index().reset_index(drop=True)
# 记录原索引与位置的映射
orig_index_map = df.index.copy()

2. 生成广播条件矩阵，合并两个条件：

# 获取每个t_index对应的位置
t_pos = orig_index_map.get_indexer(t_indices).reshape(-1, 1)
# 生成df所有行的位置数组
all_pos = np.arange(len(df)).reshape(1, -1)

# 条件1：位置大于t_index对应的位置
cond1 = all_pos > t_pos
# 条件2：Low列值小于等于对应low（广播匹配）
cond2 = df['Low'].values <= t_lows.reshape(-1, 1)
# 合并两个条件
combined_cond = cond1 & cond2

3. 提取首个符合条件的位置并映射回原索引：

# argmax返回每行第一个True的位置
first_match_pos = combined_cond.argmax(axis=1)
# 标记无匹配的情况（原代码会报错，这里用nan标记）
no_match_mask = ~combined_cond.any(axis=1)
first_match_pos[no_match_mask] = -1

# 映射回原索引
t_ixs = orig_index_map[first_match_pos]
t_ixs[no_match_mask] = np.nan

注意事项 #

• 必须保证df的索引单调递增，否则所有方案都会失效，无序索引需先执行df.sort_index()。
• 若存在无符合条件的行，原代码会抛出IndexError，上述方案通过掩码处理了这种情况，可根据需求调整标记方式。

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