如何用矩阵运算替代循环实现NumPy数组多维切片提速？

Hey there! Nice catch switching from a for loop to vectorized operations—this is exactly how you get the most out of NumPy's speed. Let's break down how to rewrite your code to ditch that loop entirely.

原代码回顾 #

First, let's recap what you're doing:

• You have a 4D array foo with shape (6,10,10,10)
• You need to generate 100 slices, each starting at an index from startIndices (shape (100,))
• Each slice takes 3 elements along foo's 2nd axis, 4 along the 3rd, and 5 along the 4th
• Results get stored in target with shape (100,6,3,4,5)

向量化解决方案 #

Here's the vectorized code that replaces your for loop entirely:

import numpy as np

foo = np.arange(6000).reshape(6,10,10,10)
startIndices = np.random.randint(5, size=(100))

# Expand startIndices to enable broadcasting with slice ranges
start_idx = startIndices[:, np.newaxis, np.newaxis, np.newaxis]  # Shape: (100,1,1,1)

# Generate indices for each slice dimension
dim1 = start_idx + np.arange(3)  # Shape: (100,3,1,1)
dim2 = start_idx + np.arange(4)  # Shape: (100,1,4,1)
dim3 = start_idx + np.arange(5)  # Shape: (100,1,1,5)

# Use broadcasting to slice foo across all start indices at once
target = foo[:, dim1, dim2, dim3].transpose(1, 0, 2, 3, 4)

为什么这能行？ #

Let's unpack the magic here:

1. Broadcasting: By adding np.newaxis to startIndices, we turn it into a 4D array that can "line up" with the slice range arrays (np.arange(3), etc.). This lets us generate all 100 sets of slice indices in one go, no loop needed.
2. Advanced Indexing: When we index foo[:, dim1, dim2, dim3], NumPy uses broadcasting to match the dimensions:
   • foo has shape (6,10,10,10)
   • dim1, dim2, dim3 broadcast together to shape (100,3,4,5)
   • The result of foo[:, dim1, dim2, dim3] is (6,100,3,4,5)
3. Transpose: We just swap the first two axes with .transpose(1,0,2,3,4) to get the target shape (100,6,3,4,5) that matches your original loop's output.

验证正确性 #

If you want to double-check this works the same as your original loop, you can run a quick test:

# Original loop implementation for comparison
target_loop = np.zeros((100,6,3,4,5))
for i in range(len(target_loop)):
    startIdx = startIndices[i]
    target_loop[i, :] = foo[:, startIdx:startIdx+3, startIdx:startIdx+4, startIdx:startIdx+5]

# Check if results are identical
print(np.array_equal(target, target_loop))  # Should print True

性能提升 #

This vectorized approach will be significantly faster than the for loop, especially as the size of startIndices grows. NumPy is optimized for these kinds of bulk operations, avoiding the Python loop overhead entirely.

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