在NumPy中高效筛选指定索引子集的方法

Got it, let's break down how to efficiently extract an index subset from your NumPy coords tuple based on the index tuple you have. First, let's align on the setup: you've got coords as a tuple of 4 NumPy arrays (each representing indices along a different axis, all the same length), and index is another tuple of arrays with a similar structure that defines which entries you want to keep from coords.

I'll cover the most common scenarios and the fastest NumPy-native ways to implement them—avoiding Python loops is key for efficiency here.

Scenario 1: Match Exact Tuples Between coords and index #

If you want to extract all entries in coords where the full (axis0, axis1, axis2, axis3) tuple exists in the set of tuples defined by index, use structured arrays for fast vectorized comparison:

import numpy as np

# Your example coords
coords = (
    np.asarray([0,0,0,1,1,1,1,1,2,2,2,3,3,3,3,4,4,4,5,5,5,5,5,6,6,6]),
    np.asarray([2,2,8,2,2,4,4,6,2,2,6,2,2,4,6,2,2,6,2,2,4,4,6,2,2,6]),
    np.asarray([0]*26),
    np.asarray([0,1,0,0,1,0,1,1,0,1,1,0,1,1,1,0,1,1,0,1,0,1,1,0,1,1])
)

# Example index tuple (subset of coords tuples)
index = (
    np.asarray([0,1,1,2,3,5]),
    np.asarray([2,2,4,6,4,4]),
    np.asarray([0,0,0,0,0,0]),
    np.asarray([0,1,1,1,1,1])
)

# Step 1: Convert coords and index into 2D arrays (each row = a full index tuple)
coords_2d = np.stack(coords, axis=1)  # Shape: (26, 4)
index_2d = np.stack(index, axis=1)   # Shape: (6, 4)

# Step 2: Use structured dtypes to enable fast tuple comparison
dt = np.dtype([('a', coords_2d.dtype), ('b', coords_2d.dtype), 
               ('c', coords_2d.dtype), ('d', coords_2d.dtype)])
coords_struct = coords_2d.view(dt).ravel()
index_struct = index_2d.view(dt).ravel()

# Step 3: Create mask and extract subset
match_mask = np.isin(coords_struct, index_struct)
coords_subset = tuple(arr[match_mask] for arr in coords)

This method is efficient because it leverages NumPy's optimized C-level operations for structured array comparisons, avoiding slow Python loops even for large datasets.

Scenario 2: Select Entries Where Each Axis Value Is in Corresponding index Array #

If you want to keep entries where coords[0][i] is in index[0], coords[1][i] is in index[1], etc., use vectorized boolean masking:

# Initialize mask to keep all entries
match_mask = np.ones(len(coords[0]), dtype=bool)

# Narrow down mask with per-axis checks
for coord_arr, index_arr in zip(coords, index):
    match_mask &= np.isin(coord_arr, index_arr)

# Extract subset
coords_subset = tuple(arr[match_mask] for arr in coords)

Each np.isin call is vectorized, so this is far faster than checking each element individually in Python.

Scenario 3: Element-Wise Per-Axis Matching #

If index has the same length as coords arrays, and you want to keep entries where coords[n][i] == index[n][i] for all axes n, use np.all for a concise, fast check:

# Create per-axis boolean arrays
axis_matches = [coord_arr == index_arr for coord_arr, index_arr in zip(coords, index)]

# Combine into a single mask (all axes must match)
match_mask = np.all(axis_matches, axis=0)

# Extract subset
coords_subset = tuple(arr[match_mask] for arr in coords)

Bonus: Get Matching Position Indices #

If you need the original indices (positions) of the matching entries in coords, use np.where to get integer indices directly:

matching_positions = np.where(match_mask)[0]
coords_subset = tuple(arr[matching_positions] for arr in coords)

Integer indexing is often slightly faster than boolean masking in NumPy, so this is a good option if you need the positions anyway.

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