如何高效将DataFrame中的最近值映射至NumPy数组

问题描述 #

我有一个包含code列和离散值列的DataFrame，还有一个存储实验值的NumPy数组。需要生成一个NumPy数组，其中每个元素对应DataFrame中与实验值最接近的离散值关联的code。

我习惯用Pandas，但当前通过嵌套循环实现的方案速度太慢，实际数据集是500×1500，还要处理700多组，急需高效方案。当前循环实现代码如下：

import numpy as np
import pandas as pd
from pandas import DataFrame


def main():
    npsize = (2,4)
    # 创建-0.75到0.25之间的随机数组
    data = np.random.uniform(-0.75,0.25,npsize)
    
    # 用于映射的DataFrame
    codes = [np.array([1,2,3]),np.array([4,5,6]),np.array([7,8,9]),np.array([10,11,12]),np.array([13,14,15])]
    values = [-0.75,-0.5,-0.25,0,0.25]
    d = {'code':codes, 'value':values}
    data_map = pd.DataFrame(data=d)

    # 需要将data中的每个元素关联到data_map中最接近value对应的code
    # 例如：-0.05 对应 [10,11,12] 
    
    # 低效的循环实现，肯定有更快的方法！
    mapped_data = np.zeros(shape=(2,4,3))
    xctr = 0
    yctr = 0

    while xctr < npsize[0]:
        while yctr < npsize[1]:
            nearest_code = data_map.iloc[(data_map['value']-data[xctr,yctr]).abs().argsort()[:1]].code.iloc[0]
            mapped_data[xctr,yctr] = nearest_code
            yctr = yctr + 1
        yctr = 0
        xctr = xctr + 1

    print (mapped_data)

if __name__ == "__main__":
    main()

高效实现方案 #

观察到data_map中的value是升序排列的，我们可以利用NumPy的向量化操作和二分查找特性，彻底替代低效的Python循环：

• 先将映射关系中的values和codes转为NumPy数组，方便后续向量化处理
• 使用np.searchsorted快速定位每个实验值在有序values中的插入位置，这一步是O(N log M)的时间复杂度（N为实验值数量，M为离散值数量）
• 针对插入位置的边界情况做处理，然后比较插入位置前后的离散值，选择距离更近的那个
• 最后通过索引直接映射到对应的code，再重塑为目标形状

优化后代码 #

import numpy as np
import pandas as pd

def main():
    npsize = (2,4)
    data = np.random.uniform(-0.75,0.25,npsize)
    
    # 保持你习惯的DataFrame映射格式
    codes = [np.array([1,2,3]),np.array([4,5,6]),np.array([7,8,9]),np.array([10,11,12]),np.array([13,14,15])]
    values = [-0.75,-0.5,-0.25,0,0.25]
    data_map = pd.DataFrame({'code':codes, 'value':values})

    # 核心高效处理逻辑
    sorted_vals = data_map['value'].values  # 提取有序的离散值数组
    code_arrays = np.stack(data_map['code'].values)  # 将code转为(5,3)的二维NumPy数组

    # 展平实验值数组，方便向量化操作
    flat_data = data.flatten()
    # 二分查找每个值的插入位置
    idx = np.searchsorted(sorted_vals, flat_data)

    # 处理边界：插入位置为0时只能选第一个值；插入位置等于数组长度时只能选最后一个值
    idx = np.clip(idx, 1, len(sorted_vals)-1)
    # 比较前后两个离散值与当前实验值的距离
    prev_dist = flat_data - sorted_vals[idx-1]
    next_dist = sorted_vals[idx] - flat_data
    # 确定最终要选的离散值索引
    final_idx = np.where(prev_dist <= next_dist, idx-1, idx)

    # 获取对应的code并重塑为目标形状
    mapped_data = code_arrays[final_idx].reshape(*data.shape, 3)

    print(mapped_data)

if __name__ == "__main__":
    main()

效率说明 #

• 全程使用NumPy向量化操作，避免了Python循环的巨大开销，对于500×1500的数据集，处理速度会比循环版快几十甚至上百倍
• 二分查找的特性让查找逻辑的时间复杂度大幅降低，再加上向量化运算，完全能支撑700多组数据的处理需求

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