如何加速Python中的scissor函数运行效率？

问题现状 #

scissor函数用于为神经网络生成训练数据集，但随着数据规模增长，耗时呈指数级上升：

• 947条数据（4小时）：0.002秒
• 4384条数据（1小时）：0.05秒
• 18132条数据（15分钟）：0.98秒
• 54793条数据（5分钟）：19.67秒
• 274761条数据：668.7秒

目前处于测试阶段，需要多次运行该函数，百万级数据的耗时完全无法接受。代码功能正常，核心需求是优化运行效率，同时排查潜在问题。

原代码 #

def scissor(scaled_data, scaled_dataframe, k):
    target_index = 20
    columns = scaled_dataframe.columns
    used_starting_points = []
    x_train = []
    y_train = []
    y_train_index = []

    available_starting_points = set(range(2 * k - 1, len(scaled_data) - (k - 1)))

    while available_starting_points:
        available_starting_points = set(range(2 * k - 1, len(scaled_data) - (k - 1))) - set(used_starting_points)
        if not available_starting_points:
            break
        starting_point = random.choice(list(available_starting_points))
        window = scaled_data[starting_point - k: starting_point]

        target_value = scaled_data[starting_point, target_index]
        y_train_index.append(scaled_dataframe.index[starting_point + (k - 1)])

        x_train.append(window)
        y_train.append(target_value)

        used_starting_points.extend(range(starting_point - k, starting_point + k))
        overlapping_points = range(starting_point - k, starting_point + k)
        available_starting_points -= set(overlapping_points)

    x_train = np.array(x_train).reshape(len(x_train), k * len(columns))
    y_train = np.array(y_train).reshape(len(y_train), )

    return x_train, y_train, y_train_index

性能瓶颈分析 #

1. 重复创建集合：每次循环都重新生成整个可用起始点集合，大数据集下集合的创建与减法操作开销极大
2. 列表转集合频繁操作：used_starting_points是列表，每次转集合都要遍历全部元素，数据量越大越慢
3. 低效随机选择：每次将集合转成列表再做随机选择，额外增加了类型转换开销
4. 边界隐患：range(starting_point - k, starting_point + k)可能超出数据集索引范围，存在越界风险

优化后的代码 #

import numpy as np
import random

def scissor_optimized(scaled_data, scaled_dataframe, k):
    target_index = 20
    num_columns = len(scaled_dataframe.columns)
    total_data = len(scaled_data)
    
    # 用布尔数组标记可用位置，True表示可用
    start_min = 2 * k - 1
    start_max = total_data - (k - 1)
    available = np.ones(total_data, dtype=bool)
    available[:start_min] = False
    available[start_max:] = False
    
    x_train = []
    y_train = []
    y_train_index = []
    
    # 获取初始可用起始点列表
    available_starts = np.where(available[start_min:start_max])[0] + start_min
    
    while len(available_starts) > 0:
        # 随机选择起始点
        idx = random.randint(0, len(available_starts) - 1)
        starting_point = available_starts[idx]
        
        # 提取窗口与目标值
        window = scaled_data[starting_point - k : starting_point]
        target_value = scaled_data[starting_point, target_index]
        y_train_index.append(scaled_dataframe.index[starting_point + (k - 1)])
        
        x_train.append(window)
        y_train.append(target_value)
        
        # 标记重叠区域为不可用，处理边界避免越界
        overlap_start = max(starting_point - k, 0)
        overlap_end = min(starting_point + k, total_data)
        available[overlap_start:overlap_end] = False
        
        # 更新可用起始点列表
        available_starts = np.where(available[start_min:start_max])[0] + start_min
    
    # 转换为numpy数组并重塑形状
    x_train = np.array(x_train).reshape(len(x_train), k * num_columns)
    y_train = np.array(y_train).flatten()
    
    return x_train, y_train, y_train_index

优化说明 #

• 用numpy布尔数组替代集合：数组操作比集合快几个数量级，尤其适合大数据集的标记与查询
• 避免重复计算：仅在初始化时设置可用范围，后续通过更新数组快速获取可用起始点
• 高效随机选择：直接在可用起始点的数组索引上随机选择，省去集合转列表的开销
• 修复边界问题：用max和min限制重叠区域范围，避免索引越界
• 精简变量：去掉不必要的变量引用，减少内存占用

额外提速建议 #

• 预分配数组：如果能估算最终训练集的大小，可直接预分配numpy数组，避免列表追加后的内存拷贝
• 分块处理：超大数据集可考虑分块并行处理，注意保证起始点不重叠的逻辑一致性

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