如何基于指定维度的2D布尔矩阵对3D NumPy数组进行元素选择？

Got it, let's work through this indexing problem together! Your data array has shape (3, 4, 3) (dimensions: i0, i1, i2), and you want to select all elements data[i0, i1, i2] where select[i0, i2] == True—where select is a (3, 3) boolean matrix matching the i0 and i2 dimensions.

Method 1: Broadcast the Boolean Mask (Most Concise) #

The key here is to expand your select mask to match the dimensionality of data, so numpy can broadcast it across the i1 dimension (the middle one, size 4). You can add a new axis to select using [:, None, :] (or [:, np.newaxis, :]) to turn its shape from (3,3) to (3,1,3)—this aligns perfectly with data's (3,4,3) shape.

Here's the code:

import numpy as np

# Your original data
data = np.arange(36).reshape((3, 4, 3))
select = np.random.choice([0, 1], size=9).reshape((3, 3)).astype(bool)

# Expand the select mask to match data's dimensions
select_expanded = select[:, None, :]

# Apply the mask directly to data
selected_elements = data[select_expanded]

What this does: Numpy automatically broadcasts the (3,1,3) mask to (3,4,3), so every i1 position for a valid (i0,i2) pair gets included in the result. The output will be a 1D array of all matching elements.

Method 2: Explicitly Get Indices (More Intuitive) #

If you prefer a more step-by-step approach, you can first find all the (i0,i2) pairs where select is True, then extract all i1 elements for those pairs:

# Get all (i0, i2) indices where select is True
i0_indices, i2_indices = np.where(select)

# Extract all i1 elements for each valid (i0,i2) pair, then flatten
selected_elements = data[i0_indices, :, i2_indices].flatten()

This gives exactly the same result as Method 1. The data[i0_indices, :, i2_indices] creates a 2D array of shape (number_of_true_pairs, 4), and flatten() turns it into a 1D array of all selected elements.

Quick Verification #

Let's say select[0,0] is True—both methods will include all elements data[0, 0, 0], data[0,1,0], data[0,2,0], data[0,3,0] in the result, which is exactly what you need.

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