如何仅在mask指定索引处计算numpy数组，避免无效运算？

核心思路 #

直接跳过mask标记的无效位置，只对有效索引执行data[i,j] * data[k,l].conj()计算，彻底避免全量O(N^4)的冗余运算，只处理需要计算的元素。

方法一：基于稀疏索引的直接计算 #

针对任意形状的四维mask，提取其中有效位置的索引，仅对这些位置执行计算后回填结果：

import numpy as np

# 示例数据（替换为你的二维复数数组）
data = np.random.rand(100, 100) + 1j * np.random.rand(100, 100)
# 示例四维mask（替换为你的实际布尔掩码）
mask = np.random.rand(100, 100, 100, 100) > 0.9

# 提取mask中所有有效位置的四维索引
i, j, k, l = np.where(mask)

# 仅对有效位置计算乘积
products = data[i, j] * data[k, l].conj()

# 初始化结果数组（按需设置初始值，比如0）
result = np.zeros_like(mask, dtype=np.complex128)
# 将计算结果回填到对应索引位置
result[i, j, k, l] = products

这种方案的时间复杂度仅与mask中True的数量成正比，适合稀疏场景，完全不会计算无用值。

方法二：针对可拆分mask的优化 #

如果你的四维mask是由两个二维mask组合而来（比如mask = mask_ij[:, :, None, None] & mask_kl[None, None, :, :]），可以进一步减少索引处理开销：

# 示例二维拆分mask
mask_ij = np.random.rand(100, 100) > 0.5
mask_kl = np.random.rand(100, 100) > 0.5
mask = mask_ij[:, :, None, None] & mask_kl[None, None, :, :]

# 提取两个二维mask对应的有效元素
valid_ij = data[mask_ij]
valid_kl_conj = data[mask_kl].conj()

# 计算有效元素的外积共轭乘积
products = valid_ij[:, None] * valid_kl_conj[None, :]

# 生成对应四维索引网格
i, j = np.where(mask_ij)
k, l = np.where(mask_kl)
ii, kk = np.meshgrid(i, k, indexing='ij')
jj, ll = np.meshgrid(j, l, indexing='ij')

# 初始化并回填结果
result = np.zeros((100,100,100,100), dtype=np.complex128)
result[ii, jj, kk, ll] = products

这种方式利用了mask的结构规律，减少了索引提取的次数，效率比通用方案更高。

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