如何用Numpy方法优化数组移动平均计算函数？

Hey there! Your current moving average function does exactly what you need, but since NumPy is built to shine with vectorized operations (no slow Python loops), we can make this way more efficient—let’s break it down.

First, let’s recap your logic #

You’re calculating the average of each pair of adjacent rows, column by column (a sliding window of size 2, step 1, applied vertically to every column). Your nested loops work, but they’re not leveraging NumPy’s core strengths.

The Fast, Vectorized Solution #

Instead of looping through each row and column, we can directly slice the array to grab all "current" rows and all "next" rows, then compute the average in one go. Here’s how:

import numpy as np

# Your sample data
b = np.array([[1, 2, 3], [2, 3, 4], [3, 4, 5], [6, 7, 8], [4, 5, 6]])

def fast_movAvg(arr):
    # Slice arr to get rows 0 to n-2, and rows 1 to n-1, then average
    return (arr[:-1] + arr[1:]) * 0.5

# Test it out
result = fast_movAvg(b)
print(result)

Output (matches your original result):

[[1.5 2.5 3.5]
 [2.5 3.5 4.5]
 [4.5 5.5 6.5]
 [5.  6.  7. ]]

Why this is better: #

• Speed: NumPy’s vectorized operations run in optimized C code, not slow Python loops. For large arrays (e.g., 10,000 rows × 100 columns), this will be 100x+ faster than your original loop-based approach.
• Readability: The code is concise and directly expresses the "average adjacent rows" logic without nested loops.
• Scalability: This approach works seamlessly for any number of rows/columns, no need to adjust loop ranges.

Bonus: Generalizing to Larger Window Sizes #

If you ever need to use a sliding window larger than 2, you can use NumPy’s sliding_window_view (available in NumPy 1.20+) for a flexible, still fully vectorized solution:

def general_movAvg(arr, window_size=2):
    # Create sliding windows along the row axis
    windows = np.lib.stride_tricks.sliding_window_view(arr, window_shape=(window_size, arr.shape[1]), axis=0)
    # Compute mean across the window axis
    return np.mean(windows, axis=1)

# Example with window size 3
print(general_movAvg(b, window_size=3))

Quick Performance Check #

Let’s use timeit to see the difference with a large array:

import timeit

# Create a large test array
large_arr = np.random.rand(10000, 100)

# Time original function
original_time = timeit.timeit(lambda: mod_movAvg(large_arr), number=100)
# Time fast function
fast_time = timeit.timeit(lambda: fast_movAvg(large_arr), number=100)

print(f"Original loop time: {original_time:.2f} seconds")
print(f"Fast vectorized time: {fast_time:.4f} seconds")

On most systems, you’ll see the vectorized version finish in milliseconds while the loop-based one takes seconds.

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