如何加速Python中该大数据集匹配运算的执行速度？

Hey there! Let's break this down for you first: that 21-minute runtime is technically "expected" given how your code is structured, but we can absolutely speed this up by orders of magnitude. Here's why and how:

Why your original code is so slow #

You’ve got two major efficiency bottlenecks here:

• Massive redundant computation: With 72,805 entries in dset1 and 1.45 million in dset2, your loop runs ~1.06 trillion individual calculations (72805 × 1455873). That’s an astronomical number of operations for a Python loop.
• np.append overhead: Every time you call np.append, it creates a brand new array and copies all existing data into it. As flag grows, this copy operation gets slower and slower, adding unnecessary drag.

The fastest fix: Use Numpy vectorization #

Numpy was built for exactly this kind of problem—its vectorized operations run optimized C code under the hood, skipping Python loop overhead entirely. Here’s the optimized version:

# Extract the first columns of both datasets
col1 = dset1[:, 0].reshape(-1, 1)  # Shape: (72805, 1) (column vector)
col2 = dset2[:, 0].reshape(1, -1)  # Shape: (1, 1455873) (row vector)

# Use broadcasting to compute all comparisons at once, then check for matches
flag = np.any(np.abs(col1 - col2) / 1000 <= 500, axis=1)

How this works: #

• By reshaping col1 and col2, Numpy automatically "broadcasts" them into a 2D matrix of shape (72805, 1455873), doing all subtraction operations in one go.
• np.any(..., axis=1) checks each row (corresponding to one entry in dset1) for any matching entry in dset2, giving you the final flag array directly.

This should run in tens of seconds to a few minutes (depending on your hardware) instead of 21 minutes. Note: This uses ~5.3GB of memory for float64 data—if your machine can’t handle that, try the block approach below.

If memory is tight: Process in blocks #

If full vectorization uses too much RAM, split dset1 into smaller blocks to reduce memory pressure:

flag = np.array([], dtype=np.bool_)
block_size = 1000  # Adjust based on your available memory

for i in range(0, len(dset1), block_size):
    # Process a chunk of dset1 at a time
    col1_block = dset1[i:i+block_size, 0].reshape(-1, 1)
    block_flag = np.any(np.abs(col1_block - dset2[:, 0]) / 1000 <= 500, axis=1)
    flag = np.concatenate([flag, block_flag])

This is slower than full vectorization but still way faster than your original loop.

Bonus: Speed up loops with Numba #

If you prefer to keep a loop structure (e.g., for extreme memory constraints), use numba to compile the loop into optimized machine code:

from numba import jit

@jit(nopython=True)
def compute_flag(dset1_col, dset2_col):
    flag = np.zeros(len(dset1_col), dtype=np.bool_)
    for i in range(len(dset1_col)):
        # Break early once we find a match (saves computation!)
        for j in range(len(dset2_col)):
            if abs(dset1_col[i] - dset2_col[j]) / 1000 <= 500:
                flag[i] = True
                break
    return flag

# Run the compiled function
flag = compute_flag(dset1[:, 0], dset2[:, 0])

The break statement here is key—it stops checking once a match is found, cutting down on unnecessary work.

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