Numpy大二维数组子矩阵筛选迭代的性能优化求助

Hey there! Let's break down why your current code is slow and fix it with some targeted optimizations.

The Root of the Slowdown #

Your existing approaches (both ndenumerate and nested loops) have a critical bottleneck: every iteration re-scans the entire array with numpy.where(M != val). For an R×C array, this means you're doing O(R×C) full array scans—leading to an overall time complexity of O(R²×C²), which gets painfully slow for large matrices. Additionally, your sub_matrix_dim function creates an empty array just to get its size, which is unnecessary overhead.

Step-by-Step Optimizations #

Let's fix this with preprocessing and vectorized operations (Numpy's superpower):

1. Preprocess Coordinates by Value

First, we'll scan the array once to group all coordinates by their value. This way, we don't have to re-scan the array every time we need coordinates of a different value.

import numpy as np
from collections import defaultdict

def preprocess_coords(M):
    value_coords = defaultdict(list)
    for idx, val in np.ndenumerate(M):
        value_coords[val].append(idx)
    # Convert lists to numpy arrays for vectorized operations later
    for val in value_coords:
        value_coords[val] = np.array(value_coords[val])
    return value_coords

2. Vectorized Filtering of Valid Coordinates

Instead of using Python loops and filter to check each coordinate individually, we'll use Numpy's vectorized operations to compute valid coordinates in bulk. This leverages Numpy's optimized C backend to avoid slow Python-level loops.

We'll also replace the inefficient sub_matrix_dim function with a direct calculation (no need to create an empty array!).

def find_valid_submatrices(M, L, H):
    R, C = M.shape
    value_coords = preprocess_coords(M)
    poss = {}
    
    for x in range(R):
        for y in range(C):
            current_val = M[x, y]
            # Gather all coordinates that don't match the current value
            all_other_coords = []
            for val, coords in value_coords.items():
                if val != current_val:
                    all_other_coords.append(coords)
            
            if not all_other_coords:
                poss[(x, y)] = []
                continue
            
            # Combine into a single numpy array for vectorized processing
            all_other_coords = np.vstack(all_other_coords)
            
            # Optional: Spatial pre-filter to reduce the number of coordinates to check
            # If H is small, we can skip coordinates that can't possibly form a valid submatrix
            max_offset = H - 1  # Since (abs(a-x)+1) * (abs(b-y)+1) <= H implies each dimension <= H
            x_bounds = (x - max_offset, x + max_offset)
            y_bounds = (y - max_offset, y + max_offset)
            spatial_mask = (all_other_coords[:, 0] >= x_bounds[0]) & \
                           (all_other_coords[:, 0] <= x_bounds[1]) & \
                           (all_other_coords[:, 1] >= y_bounds[0]) & \
                           (all_other_coords[:, 1] <= y_bounds[1])
            all_other_coords = all_other_coords[spatial_mask]
            
            # Calculate submatrix sizes in bulk
            dx = np.abs(all_other_coords[:, 0] - x) + 1
            dy = np.abs(all_other_coords[:, 1] - y) + 1
            sizes = dx * dy
            
            # Filter coordinates that meet the size constraints
            valid_mask = (sizes >= 2 * L) & (sizes <= H)
            valid_coords = all_other_coords[valid_mask]
            
            # Convert back to tuples for the dictionary
            poss[(x, y)] = [tuple(coord) for coord in valid_coords]
    
    return poss

Testing with Your Example #

Let's verify this works with your sample matrix:

M = np.array([
    [0, 0, 0, 0, 0],
    [0, 1, 1, 1, 0],
    [0, 0, 0, 0, 0]
])
L = 1
H = 6

result = find_valid_submatrices(M, L, H)
print(result[(0, 0)])  # Output: [(1, 1), (1, 2)]
print(result[(0, 1)])  # Output: [(1, 1), (1, 2), (1, 3)]

Why This Is Faster #

• Preprocessing: We only scan the array once to group coordinates, instead of O(R×C) times.
• Vectorization: Bulk calculations with Numpy are orders of magnitude faster than Python loops and filter.
• Spatial Pre-Filter: Reduces the number of coordinates we need to check, especially when H is small compared to the array size.
• Removed Overhead: No more empty array creation just to get a submatrix size.

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