寻求比numpy.where更节省内存的大数组切片方案

Great question—dealing with large NumPy arrays and memory constraints is super common, so let's break this down step by step.

First, to answer your core question: Yes, multiple & operators do create repeated intermediate copies! Every time you combine two boolean conditions with &, NumPy generates a brand-new boolean array to store the intermediate result. For example, (x < threshold) & (y > threshold) first creates a boolean array for x < threshold, another for y > threshold, then a third array for their bitwise AND. For a 5 million-element array, each boolean array takes ~5MB (since bool is 1 byte per element), which might seem small, but stacked with multiple conditions + loop iterations, these intermediate arrays can pile up and trigger a MemoryError if your system is tight on RAM.

Now, here are several memory-friendly alternatives that either reduce overhead or let you retain reusable indices:

1. Use np.logical_and.reduce to minimize intermediate arrays #

Pack all your conditions into a list and combine them in one go with reduce—this only generates a single final boolean mask instead of multiple intermediate arrays:

import numpy as np

for threshold in np.arange(0, 1, 0.1):
    x = np.random.random(5000000)
    y = np.random.random(5000000)
    z = np.random.random(5000000)
    
    # Bundle all conditions into a list
    conditions = [
        x < threshold,
        y > threshold,
        z > threshold,
        z < threshold + 0.1
    ]
    # Combine all conditions in one step to create a single mask
    mask = np.logical_and.reduce(conditions)
    # Get indices if you need to reuse them later
    inds = mask.nonzero()[0]  # Identical to np.where(mask)
    
    DoSomeJunk(a[inds], b[inds], c[inds])

2. Skip np.where entirely—use boolean masks directly #

If you only need the indices to extract elements from other arrays, you don't have to generate an explicit inds array at all. Boolean masks work directly for indexing, saving you the memory needed to store millions of integers:

for threshold in np.arange(0, 1, 0.1):
    x = np.random.random(5000000)
    y = np.random.random(5000000)
    z = np.random.random(5000000)
    
    mask = (x < threshold) & (y > threshold) & (z > threshold) & (z < threshold + 0.1)
    # Use the mask directly for indexing—no inds array required
    DoSomeJunk(a[mask], b[mask], c[mask])

3. Optimize condition logic + manually free memory #

First, combine redundant conditions (like your two checks on z) to cut down on intermediate arrays. Then, explicitly delete unused variables and trigger garbage collection to prevent memory buildup during loops:

import gc

for threshold in np.arange(0, 1, 0.1):
    x = np.random.random(5000000)
    y = np.random.random(5000000)
    z = np.random.random(5000000)
    
    # Precompute z's combined condition first to reduce intermediate arrays
    z_mask = (z > threshold) & (z < threshold + 0.1)
    mask = (x < threshold) & (y > threshold) & z_mask
    inds = mask.nonzero()[0]
    
    DoSomeJunk(a[inds], b[inds], c[inds])
    
    # Manually clean up unused variables to free memory immediately
    del x, y, z, z_mask, mask, inds
    gc.collect()

4. Process the array in chunks for extreme memory constraints #

If your array is even larger than 5 million elements, split it into smaller chunks and process each chunk individually. This keeps memory usage low, as you only ever load a small portion of the array into RAM at once:

chunk_size = 1000000  # Adjust based on your system's available RAM

for threshold in np.arange(0, 1, 0.1):
    x = np.random.random(5000000)
    y = np.random.random(5000000)
    z = np.random.random(5000000)
    
    for i in range(0, len(x), chunk_size):
        # Slice the current chunk from the full arrays
        x_chunk = x[i:i+chunk_size]
        y_chunk = y[i:i+chunk_size]
        z_chunk = z[i:i+chunk_size]
        
        # Compute mask for the chunk
        mask_chunk = (x_chunk < threshold) & (y_chunk > threshold) & (z_chunk > threshold) & (z_chunk < threshold + 0.1)
        # Get global indices for the chunk
        inds_chunk = i + np.where(mask_chunk)[0]
        
        # Process the chunk's elements
        DoSomeJunk(a[inds_chunk], b[inds_chunk], c[inds_chunk])

A quick side note: np.where(mask) is essentially just a wrapper for mask.nonzero(), so they perform identically in terms of memory and speed. The real win comes from optimizing how you create the boolean mask itself.

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