如何高效实现多次将2D向量添加至目标矩阵的操作？

Hey there! Let's fix that slow padding operation you're dealing with—especially since your groups list has 122,000 entries, the repeated np.vstack calls in your loop are killing performance by forcing constant memory reallocation and data copying. Here's a much faster approach that pre-allocates memory and uses numpy's efficient slicing instead of incremental stacking.

The Core Problem with Your Current Code #

Every time you call np.vstack([padded_dataset, subtree]) or np.vstack([padded_dataset, pads]), numpy has to:

• Create a brand new array with enough space for the existing data plus the new rows
• Copy all the existing data into this new array
• Add the new rows
• Discard the old array

With 122k iterations, this adds up to an enormous amount of unnecessary work.

Optimized Implementation #

The key is to pre-calculate the total size of your final padded_dataset and allocate all memory upfront, then fill in the data directly using slicing. This eliminates all the repeated copying from vstack.

import numpy as np

# dummy function just for testing
def get_max_subtree_length(groups):
    return 20

def pad_groups_optimized(dataset, groups):
    dataset_np = np.array(dataset)
    max_subtree_length = get_max_subtree_length(groups)
    row_length = dataset_np.shape[1]
    dataset_row_length_no_first = row_length - 1
    zeros_pad = np.zeros(dataset_row_length_no_first, dtype=dataset_np.dtype)
    
    # Precompute all necessary lengths and values
    group_values = np.array([g[0] for g in groups])
    subtree_lengths = np.array([g[1] for g in groups])
    pad_counts = max_subtree_length - subtree_lengths
    
    # Calculate total rows needed and pre-allocate the output array
    total_rows = np.sum(subtree_lengths + pad_counts)
    padded_dataset = np.zeros((total_rows, row_length), dtype=dataset_np.dtype)
    
    # Track positions for filling data
    current_pos = 0
    subtree_start = 0
    
    for group_val, subtree_len, pad_count in zip(group_values, subtree_lengths, pad_counts):
        # Fill in the subtree from the original dataset
        subtree_end = current_pos + subtree_len
        padded_dataset[current_pos:subtree_end] = dataset_np[subtree_start:subtree_start+subtree_len]
        
        # Fill in the padding rows (if needed)
        if pad_count > 0:
            pad_row = np.concatenate([[group_val], zeros_pad])
            pad_end = subtree_end + pad_count
            # Numpy efficiently broadcasts the single pad row to all required positions
            padded_dataset[subtree_end:pad_end] = pad_row
        
        # Update positions for next iteration
        current_pos = pad_end if pad_count > 0 else subtree_end
        subtree_start += subtree_len
    
    return padded_dataset

Why This Is Faster #

1. Pre-allocated Memory: We calculate exactly how big the final array needs to be and create it once—no more repeated memory allocation and copying.
2. Efficient Slicing: Numpy's slice assignments are optimized at the C level, way faster than looping and stacking small arrays.
3. Minimal Temporary Arrays: We only create the pad row once per group instead of generating a full pads array with np.repeat, saving memory and processing time.
4. Type Consistency: We use the same dtype as the original dataset to avoid unnecessary type conversions.

Testing the Optimized Code #

Using your test example, this will produce the exact same output as your original function, but in a fraction of the time:

# Test the optimized function
dataset = np.array([
 [1, 2, 3],
 [1, 2, 3],
 [1, 2, 3],
 [2, 2, 3],
 [2, 2, 3],
 [2, 2, 3],
 [3, 2, 3],
 [3, 2, 3],
 [3, 2, 3],
 [4, 2, 3],
 [4, 2, 3],
 [4, 2, 3],
 [5, 2, 3],
 [5, 2, 3],
 [5, 2, 3],
 [6, 2, 3],
 [6, 2, 3],
 [6, 2, 3],
 [7, 2, 3],
 [7, 2, 3],
 [7, 2, 3],
 [8, 2, 3],
 [8, 2, 3],
 [8, 2, 3]])
groups = [(1, 3), (2, 3), (3, 3), (4, 3), (5, 3), (6, 3), (7, 3), (8, 3)]

padded_dataset_opt = pad_groups_optimized(dataset, groups)
print(len(padded_dataset_opt))  # Output: 160
# The content will match your original output exactly

Additional Performance Tips #

• If get_max_subtree_length is expensive, keep calling it only once (your original code already does this, which is great!).
• If your input dataset is already a numpy array, skip the np.array(dataset) conversion to save a little time.
• For extremely large datasets that push memory limits, consider using np.memmap to work with data on disk, but pre-allocation should handle 122k groups easily.

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