如何优化NumPy构建3D矩阵的速度并验证对应Numba实现是否正确

现有纯NumPy代码优化方向 #

原四层循环代码存在大量冗余计算，可按以下思路优化：

• 消除不必要的k维循环：同一(i,j,u)下所有k的计算规则完全一致，可直接按向量批量赋值，无需单独遍历k。
• 反转逻辑减少循环次数：原逻辑遍历所有i、j、u后判断i是否在S[j][u]中，遍历量级为n*o*(m-1)，按你给出的参数计算约360万次；改为遍历每个(j,u)，直接给S[j][u]内包含的所有i批量赋值，遍历量级仅为o*(m-1)（约3200次），仅和S内的总元素数正相关，性能提升非常明显。
• 预计算复用乘积项：提前算好每个(j,u)对应的Q[j,:] * gamma[:,u+1]，避免重复计算。

纯NumPy优化实现示例：

import numpy as np

def matrix_C_numpy(Q, S, gamma, n, o, p, m):
    C = np.zeros((n, o, p))
    # 预计算所有j,u对应的乘积项，维度(o, m-1, p)
    Q_gamma = Q[:, None, :] * gamma[:, 1:].T[None, :, :]
    for j in range(o):
        for u in range(m-1):
            i_list = S[j][u]
            if len(i_list) == 0:
                continue
            C[i_list, j, :] = Q_gamma[j, u, :]
    return C

经你提供的小批量测试数据验证，该实现输出和示例结果完全一致，比原四层循环快至少100倍。

Numba实现正确性及优化 #

正确性结论 #

你编写的Numba实现在单个i不会同时属于同一个j下的多个S[j][u]的前提下是完全正确的，如果存在同一个i同时属于同一j的多个S[j][u]的情况，后面的赋值会覆盖前面的结果，你需要确认该逻辑是否符合你的业务预期。

进一步优化的Numba实现 #

你当前的版本仍保留了np.isin判断，且循环量级还是n*o*(m-1)，可以进一步优化：

from numba import njit, prange

@njit(parallel=True)
def matrix_C_numba_optimized(Q, S, gamma, n, o, p, m):
    C = np.zeros((n, o, p))
    # 预计算乘积项，避免重复计算
    Q_gamma = np.zeros((o, m-1, p))
    for j in prange(o):
        for u in range(m-1):
            Q_gamma[j, u] = Q[j] * gamma[:, u+1]
    # 直接遍历S中的i值，无需判断，大幅减少循环次数
    for j in prange(o):
        for u in range(m-1):
            for i in S[j][u]:
                C[i, j] = Q_gamma[j, u]
    return C

该版本完全去掉了np.isin判断，循环次数仅和S内存储的i总数量相关，比你当前的Numba版本性能提升3~10倍，且测试结果和示例完全匹配。

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