Pandas优化：多条件DataFrame反向嵌套循环的矢量化改造

Your original code works correctly, but the nested loops and repeated slice-based mean calculations make it extremely slow when called tens of thousands of times. The core issue is that Python loops are inherently slow for large datasets, and slicing DataFrames to compute means repeatedly incurs significant overhead (each slice operation and mean calculation is O(k), where k is the length of the interval).

Let's fix this by leveraging cumulative sums (cumsum) for fast interval mean calculations and NumPy vectorization to replace inner loops, which will drastically speed up your code.

Key Optimizations Explained #

1. Cumulative Sum for Fast Mean Calculation: Instead of slicing the DataFrame to compute means every time, precompute cumulative sums for the volume columns. This lets you calculate the mean of any interval [start, end) in O(1) time using:

   interval_mean = (cumulative_sum[end] - cumulative_sum[start]) / (end - start)
   

2. Vectorized Inner Loop: Replace the inner Python loop with NumPy array operations to batch-process all candidate end positions for a given start position. This eliminates the overhead of Python loop iteration.
3. Retain Early Termination: Keep your original prev_end_pos optimization to reduce the number of candidate end positions for subsequent start positions, as we only need the largest valid interval for the first qualifying start position.

Optimized Code #

import pandas as pd
import numpy as np

def optimizedFunction():
    # Configuration variables
    minVolume = 2000
    exchange1 = 'binance'
    exchange2 = 'bitmart'
    volEx1Str = 'volume_' + exchange1
    volEx2Str = 'volume_' + exchange2
    threshold = 15.0
    minDuration = 10.0

    # Load dataset
    dataset = pd.read_csv('example.csv', sep='|')
    total_rows = len(dataset)

    # Get iloc positions of rows where diffprice meets threshold
    # Convert original index to 0-based iloc positions
    base_index = dataset.index[0]
    indices_thresh = (dataset.index[dataset.diffprice >= threshold].values - base_index).astype(int)

    # Precompute cumulative sums for volume columns (add leading 0 for interval calculations)
    sum_vol1 = np.concatenate([[0], dataset[volEx1Str].cumsum().values])
    sum_vol2 = np.concatenate([[0], dataset[volEx2Str].cumsum().values])

    prev_end_pos = total_rows
    pv = None

    for start_pos in indices_thresh:
        # Calculate minimum end position to meet duration requirement
        min_end_pos = start_pos + int(minDuration)
        if min_end_pos >= prev_end_pos:
            continue  # Skip if even the smallest valid interval is too short

        # Generate all candidate end positions in descending order (from largest to smallest)
        end_pos_candidates = np.arange(prev_end_pos, min_end_pos - 1, -1)
        interval_lengths = end_pos_candidates - start_pos

        # Vectorized calculation of average volumes for all candidate intervals
        sum1_diff = sum_vol1[end_pos_candidates] - sum_vol1[start_pos]
        avg_vol1 = sum1_diff / interval_lengths
        sum2_diff = sum_vol2[end_pos_candidates] - sum_vol2[start_pos]
        avg_vol2 = sum2_diff / interval_lengths

        # Find the first (largest) valid end position
        valid_mask = (avg_vol1 > minVolume) & (avg_vol2 > minVolume)
        valid_indices = np.where(valid_mask)[0]

        if len(valid_indices) > 0:
            selected_end_pos = end_pos_candidates[valid_indices[0]]
            pv = dataset.iloc[start_pos:selected_end_pos].copy()
            break  # Exit early once we find the first valid interval

        # Update prev_end_pos to reduce candidates for next start positions
        prev_end_pos = start_pos

    if pv is None:
        print('No combination could be found for this iteration.')
        return
    return pv

Why This Is Faster #

• O(1) Mean Calculations: Cumulative sums let us compute interval means in constant time instead of linear time for each slice.
• Vectorized Operations: NumPy handles the inner loop calculations in optimized C code, which is orders of magnitude faster than pure Python loops.
• Reduced Candidate Pool: The prev_end_pos variable ensures we don't waste time checking end positions that can't yield a larger interval than previous iterations.

This code will produce the exact same output as your original function but will run significantly faster—especially when called tens of thousands of times. For even larger datasets, you could explore further optimizations like precomputing all valid intervals upfront, but this implementation strikes a great balance between speed and readability.

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