如何通过矩阵运算加速带重复索引的2D NumPy数组叉积索引更新

你之前的实现方案存在两个问题：

1. 输出数组s的初始化维度错误，双重循环逻辑中输出的维度是宽高各除以grid_size向上取整的结果，而非和输入image同维度
2. 直接使用高级索引的+=操作遇到重复索引时不会累加，只会保留最后一次写入的值，这是NumPy高级索引的默认行为

你需要的功能可以通过NumPy内置的np.add.at实现，该方法专门处理重复索引的原地累加需求，全程无显式Python循环，完全通过底层运算加速：

import numpy as np

def get_s(image, grid_size):
    W, H = image.shape
    # 计算输出矩阵的正确维度
    out_width = W // grid_size + (1 if W % grid_size != 0 else 0)
    out_height = H // grid_size + (1 if H % grid_size != 0 else 0)
    s = np.zeros((out_width, out_height), dtype=image.dtype)
    
    # 生成每个像素对应的网格坐标
    x_indices = (np.arange(W) // grid_size).astype(int)
    y_indices = (np.arange(H) // grid_size).astype(int)
    # 结合ix_构造叉积索引，add.at自动处理重复索引累加
    np.add.at(s, np.ix_(x_indices, y_indices), image)
    return s

使用ix_可以快速构造用于叉积索引的索引数组，a[np.ix_([1,3],[2,5])]会返回数组[[a[1,2] a[1,5]], [a[3,2] a[3,5]]]

如果你的输入图像宽高恰好是grid_size的整数倍，可以采用性能更高的分块求和方案，完全避免索引操作，纯矩阵运算速度提升明显：

def get_s_fast(image, grid_size):
    W, H = image.shape
    # 仅适用于宽高能被grid_size整除的场景
    assert W % grid_size == 0 and H % grid_size == 0
    # 拆分维度后对块内维度求和
    return image.reshape(
        W // grid_size, grid_size,
        H // grid_size, grid_size
    ).sum(axis=(1, 3))

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