Numpy无需创建临时数组按距离条件替换多维数组元素的方法

方案1：利用隐式广播省去显式r_temp创建（改动最小） #

你当前使用的np.broadcast_to本身生成的是视图不会占用额外内存，真正的额外内存开销来自比较时生成的布尔掩码数组，你可以直接通过调整数组形状触发numpy的隐式广播，完全省去显式创建r_temp的步骤，代码更简洁，内存占用更低：

# 给r补全前导空维度，适配r0的维度，仅修改数组视图不复制数据
r_aligned = r.reshape((1,) * (Data.ndim - 1) + (-1,))
# 直接利用广播规则做比较，无需生成全尺寸的r_temp
Data[r_aligned >= r0[..., None]] = 0

该方案和你原有实现的运算速度完全一致，仅省去了冗余的r_temp变量定义，内存开销减少了一个和Data同尺寸的临时数组占用。

方案2：基于有序r的超低内存优化（适合超大规模Data） #

如果你的距离向量r是单调递增的（绝大多数距离场景下都满足该条件），可以通过预计算切割索引的方式进一步降低内存开销。
首先预计算每个r0对应的切割位置：

# 查找每个r0元素对应的切割位置：r中第一个>=r0[i,j]的索引
# 时间复杂度为O(M*logN)，M是r0的元素总数，N是r的长度
cut_indices = np.searchsorted(r, r0)

如果内存余量允许，你可以直接生成掩码快速置零：

last_axis = Data.ndim - 1
mask = np.arange(Data.shape[last_axis]) < cut_indices[..., None]
Data *= mask

如果内存非常紧张，连全尺寸的布尔掩码都无法容纳，可以配合numba的JIT编译实现循环加速，内存开销仅为r0的大小，JIT编译后的循环速度和原生矢量运算基本持平：

from numba import jit

@jit(nopython=True)
def set_zero_by_r0(Data, cut_indices):
    # 可根据Data的前导维度数量调整循环层级，此处以三维Data为例
    for i in range(Data.shape[0]):
        for j in range(Data.shape[1]):
            k0 = cut_indices[i,j]
            Data[i,j, k0:] = 0

set_zero_by_r0(Data, cut_indices)

该方案在Data尺寸极大时，内存和速度表现都远优于原生广播方案。

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