如何无需遍历numpy数组即可在pandas匹配场景下获得相同输出

如下代码包含numpy数组clusters，其元素通过np.where与pandas DataFrame的每一行做比对，SoFunc函数接收clusters数组元素作为入参，返回所有判定条件为True的行。

目前我可以遍历该数组，将每个数组元素代入对应的np.where条件做比对。我已知遍历numpy数组效率较低，而实际使用的数据集规模会大很多，请问如何取消遍历操作仍能得到完全相同的输出？

import numpy as np
import pandas as pd

def genMockDataFrame(days,startPrice,colName,startDate,seed=None): 
    periods = days*24
    np.random.seed(seed)
    steps = np.random.normal(loc=0, scale=0.0018, size=periods)
    steps[0]=0
    P = startPrice+np.cumsum(steps)
    P = [round(i,4) for i in P]

    fxDF = pd.DataFrame({ 
        'ticker':np.repeat( [colName], periods ),
        'date':np.tile( pd.date_range(startDate, periods=periods, freq='H'), 1 ),
        'price':(P)})
    fxDF.index = pd.to_datetime(fxDF.date)
    fxDF = fxDF.price.resample('D').ohlc()
    fxDF.columns = [i.title() for i in fxDF.columns]
    return fxDF


def SoFunc(clust):
    #generate mock data
    df = genMockDataFrame(10,1.1904,'eurusd','19/3/2020',seed=157)
    df["Upper_Band"] = 1.1928
    df.loc["2020-03-27", "Upper_Band"] = 1.2118
    df.loc["2020-03-26", "Upper_Band"] = 1.2200
    df["Level"] = np.where((df["High"] >= clust)
                                      & (df["Low"] <= clust)
                                     & (df["High"] >= df["Upper_Band"] ),1,np.NaN
                                      )
    return df.dropna()

clusters = np.array([1.1929   , 1.2118 ])

l = []

for i in range(len(clusters)):
    l.append(SoFunc(clusters[i]))
    
pd.concat(l)

Open    High     Low   Close  Upper_Band  Level
date                                                         
2020-03-19  1.1904  1.1937  1.1832  1.1832      1.1928    1.0
2020-03-25  1.1939  1.1939  1.1864  1.1936      1.1928    1.0
2020-03-27  1.2118  1.2144  1.2039  1.2089      1.2118    1.0

你原来的实现有两个核心性能损耗点：一是每次遍历都重复生成完全相同的df，二是手动遍历clusters数组做判断，优化方案直接针对这两个点调整即可：

1. 把df生成、固定字段赋值的逻辑从循环里抽出来，只执行一次，避免重复计算
2. 用numpy的广播机制一次性完成所有clusters的条件匹配，全程不需要手动遍历

优化后的代码如下：

# 仅生成一次df，省去重复生成的开销
df = genMockDataFrame(10,1.1904,'eurusd','19/3/2020',seed=157)
df["Upper_Band"] = 1.1928
df.loc["2020-03-27", "Upper_Band"] = 1.2118
df.loc["2020-03-26", "Upper_Band"] = 1.2200

clusters = np.array([1.1929   , 1.2118 ])

# 转成numpy数组做广播匹配，一次性完成所有clust的条件判断
high_arr = df["High"].to_numpy()[:, None]
low_arr = df["Low"].to_numpy()[:, None]
upper_cond = (df["High"] >= df["Upper_Band"]).to_numpy()[:, None]
all_cond = (high_arr >= clusters) & (low_arr <= clusters) & upper_cond

# 只要有任意一个clust满足条件就保留该行
mask = all_cond.any(axis=1)
df["Level"] = np.where(mask, 1, np.nan)
result = df.dropna()

运行后得到的result和你原来遍历输出的结果完全一致，数据量越大性能优势越明显，比手动遍历效率高几十到上百倍都很正常。

如果你的场景中存在单个行匹配多个clust的情况，且需要和原来遍历concat的输出完全一致（包括重复的行），可以把mask逻辑替换为以下代码：

# 提取所有满足条件的行索引，按匹配顺序取行，完全兼容原输出逻辑
match_rows = np.where(all_cond)[0]
result = df.iloc[match_rows].assign(Level=1)

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