Python性能优化：用Numpy数组替代itertools生成指定均值的组合

Your original code's performance bottlenecks come from precomputing all possible combinations (memory-heavy) and looping through each month to filter valid combinations (slow for large datasets). Below is a vectorized NumPy implementation that eliminates loops and avoids precomputing all combinations, drastically improving performance for large-scale simulations.

Key Optimizations: #

• Vectorized Operations: Process entire groups of months (by num_deals value) in bulk instead of iterating per month.
• Direct Valid Value Generation: Instead of filtering precomputed combinations, generate valid random values directly using arithmetic constraints.
• Sorted Output: Matches the original function's use of combinations_with_replacement by sorting each generated combination.

Optimized Code: #

import numpy as np

def deal_months_np(ccc, num_deals, min_val, max_val):
    # Convert inputs to NumPy arrays for vector operations
    ccc = np.asarray(ccc)
    num_deals = np.asarray(num_deals)
    
    # Validate input lengths
    if len(ccc) != len(num_deals):
        raise ValueError("ccc and num_deals must have the same length")
    
    # Initialize result array (object dtype to hold lists or 0)
    result = np.zeros(len(ccc), dtype=object)
    
    # Handle num_deals == 0
    mask_0 = num_deals == 0
    result[mask_0] = 0
    
    # Handle num_deals == 1
    mask_1 = num_deals == 1
    ccc_1 = ccc[mask_1]
    valid_1 = (ccc_1 >= min_val) & (ccc_1 <= max_val)
    result[mask_1] = [[x] if v else 0 for x, v in zip(ccc_1, valid_1)]
    
    # Handle num_deals == 2
    mask_2 = num_deals == 2
    ccc_2 = ccc[mask_2]
    S_2 = 2 * ccc_2
    
    # Calculate valid range for first element
    low_x1 = np.maximum(min_val, S_2 - max_val)
    high_x1 = np.minimum(max_val, S_2 - min_val)
    valid_2 = low_x1 <= high_x1
    
    # Generate random first element and compute second
    x1 = np.where(valid_2, np.random.randint(low_x1, high_x1 + 1, size=len(ccc_2)), 0)
    x2 = S_2 - x1
    
    # Sort pairs to match combinations_with_replacement order
    pairs = np.stack([np.minimum(x1, x2), np.maximum(x1, x2)], axis=1)
    result[mask_2] = [list(p) if v else 0 for p, v in zip(pairs, valid_2)]
    
    # Handle num_deals == 3
    mask_3 = num_deals == 3
    ccc_3 = ccc[mask_3]
    S_3 = 3 * ccc_3
    
    # Calculate valid range for first element
    low_x1 = np.maximum(min_val, S_3 - 2 * max_val)
    high_x1 = np.minimum(max_val, S_3 - 2 * min_val)
    valid_3 = low_x1 <= high_x1
    
    # Generate first element
    x1 = np.where(valid_3, np.random.randint(low_x1, high_x1 + 1, size=len(ccc_3)), 0)
    S_rem1 = S_3 - x1
    
    # Calculate valid range for second element
    low_x2 = np.maximum(min_val, S_rem1 - max_val)
    high_x2 = np.minimum(max_val, S_rem1 - min_val)
    valid_3 = valid_3 & (low_x2 <= high_x2)
    
    # Generate second element and compute third
    x2 = np.where(valid_3, np.random.randint(low_x2, high_x2 + 1, size=len(ccc_3)), 0)
    x3 = S_rem1 - x2
    
    # Sort triplets to match combinations_with_replacement order
    triplets = np.stack([x1, x2, x3], axis=1)
    triplets_sorted = np.sort(triplets, axis=1)
    result[mask_3] = [list(t) if v else 0 for t, v in zip(triplets_sorted, valid_3)]
    
    # Handle num_deals == 4
    mask_4 = num_deals == 4
    ccc_4 = ccc[mask_4]
    S_4 = 4 * ccc_4
    
    # Initial validity check (sum must be within 4*min to 4*max)
    valid_4 = (S_4 >= 4 * min_val) & (S_4 <= 4 * max_val)
    
    # Calculate valid range for first element
    low_x1 = np.maximum(min_val, S_4 - 3 * max_val)
    high_x1 = np.minimum(max_val, S_4 - 3 * min_val)
    valid_4 = valid_4 & (low_x1 <= high_x1)
    
    # Generate first element
    x1 = np.where(valid_4, np.random.randint(low_x1, high_x1 + 1, size=len(ccc_4)), 0)
    S_rem1 = S_4 - x1
    
    # Calculate valid range for second element
    low_x2 = np.maximum(min_val, S_rem1 - 2 * max_val)
    high_x2 = np.minimum(max_val, S_rem1 - 2 * min_val)
    valid_4 = valid_4 & (low_x2 <= high_x2)
    
    # Generate second element
    x2 = np.where(valid_4, np.random.randint(low_x2, high_x2 + 1, size=len(ccc_4)), 0)
    S_rem2 = S_rem1 - x2
    
    # Calculate valid range for third element
    low_x3 = np.maximum(min_val, S_rem2 - max_val)
    high_x3 = np.minimum(max_val, S_rem2 - min_val)
    valid_4 = valid_4 & (low_x3 <= high_x3)
    
    # Generate third element and compute fourth
    x3 = np.where(valid_4, np.random.randint(low_x3, high_x3 + 1, size=len(ccc_4)), 0)
    x4 = S_rem2 - x3
    
    # Sort quadruplets to match combinations_with_replacement order
    quadruplets = np.stack([x1, x2, x3, x4], axis=1)
    quadruplets_sorted = np.sort(quadruplets, axis=1)
    result[mask_4] = [list(q) if v else 0 for q, v in zip(quadruplets_sorted, valid_4)]
    
    # Convert back to list to match original function's output format
    return result.tolist()

Performance Notes: #

• No Precomputed Combinations: Avoids storing thousands of combinations in memory, reducing overhead.
• Bulk Processing: Each num_deals group is processed in vectorized operations, which are orders of magnitude faster than Python loops for large datasets.
• Direct Valid Generation: Uses arithmetic constraints to generate valid values without filtering, eliminating the O(M) cost of checking each combination (where M is the number of precomputed combinations).

Usage Example: #

# Sample inputs (1000 months of data)
np.random.seed(42)
ccc = np.random.normal(loc=8, scale=3, size=1000).astype(int)
num_deals = np.random.randint(0, 5, size=1000)

# Generate results
output = deal_months_np(ccc, num_deals, min_val=0, max_val=36)

This implementation maintains the original function's behavior while delivering drastically improved performance for large-scale simulations (e.g., thousands of months of data).

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