如何用更Pythonic的方式实现带条件的numpy二维数组操作？

Instead of nested loops, you can use NumPy's vectorized operations and np.where to achieve the same result efficiently while preserving the 2D array structure. Here's the implementation:

Step-by-Step Breakdown #

1. Initialize the result array: Start with a zero array matching the shape of f.
2. Extract relevant subarrays: Isolate the inner region (excluding borders) of u and f, plus the adjacent left/right columns of f needed for calculations.
3. Apply conditional logic: Use np.where to handle both cases (positive vs non-positive u values) in a single vectorized operation.

Final Code #

import numpy as np

f = np.array([[0, 0, 0, 0, 0, 0, 0], [0, 10, 22, 30, 40, 50, 0], [0, 11, 22, 33, 44, 55, 0], [0, 0, 0, 0, 0, 0, 0]])
u = np.array([[1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, -1, 1], [1, 1, -1, -1, -1, 1, 1], [1, 1, 1, 1, 1, 1, 1]])

x = np.zeros_like(f)

# Extract inner regions and adjacent columns
u_inner = u[1:-1, 1:-1]
f_inner = f[1:-1, 1:-1]
f_left = f[1:-1, :-2]
f_right = f[1:-1, 2:]

# Apply conditional calculation using np.where
x[1:-1, 1:-1] = np.where(
    u_inner > 0,
    u_inner * (f_inner - f_left),
    (-u_inner) * (f_right - f_inner)
)

print(x)

Output Validation #

This code produces identical results to the original nested loop. For example:

• x[1,5] (where u is -1) equals -50, matching the loop's calculation.
• x[2,2] (where u is -1) equals 11, which aligns with the loop's output.

This approach avoids explicit loops, leverages NumPy's optimized backend for faster execution, and maintains the 2D structure of your arrays throughout the process.

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