如何用Numpy高效实现仅对指定j索引子集的矩阵赋值操作

场景1：allowed_j_indices 全局固定（不随i变化） #

直接对允许的j索引做切片广播，避免全量n*n矩阵的内存占用和冗余运算：

import numpy as np

n = 1000
a = np.random.rand(n)
b = np.random.rand(n)
C = np.zeros((n, n))
allowed_j_indices = [10, 20, 34, 100]

# 仅对允许的j维度做广播，运算量为 n * len(allowed_j_indices)
C[:, allowed_j_indices] = b[allowed_j_indices][None, :] - a[:, None]

该方案完全复用Numpy向量化运算的高效性，内存开销仅为全量广播的 len(allowed_j_indices)/n，当允许的j数量远小于n时优势极其明显。

场景2：allowed_j_indices 随i变化（补充需求场景） #

方案1：Numpy 高级索引（无Python层循环） #

先构造所有需要赋值的坐标对，再批量做向量化运算赋值：

import numpy as np

n = 1000
a = np.random.rand(n)
b = np.random.rand(n)
C = np.zeros((n, n))
# 示例：每个i对应的允许j索引，可按业务逻辑自定义
allowed_j_per_i = [
    [10, 20] if i % 2 == 0 else [34, 100] 
    for i in range(n)
]

# 构造坐标数组
i_coords = np.repeat(np.arange(n), [len(js) for js in allowed_j_per_i])
j_coords = np.concatenate([np.array(js) for js in allowed_j_per_i])

# 批量赋值，完全向量化运算
C[i_coords, j_coords] = b[j_coords] - a[i_coords]

运算开销仅和总待赋值元素数量成正比，无冗余运算，适合绝大多数场景。

方案2：Numba 加速循环（代码改动最小） #

如果对代码可读性要求更高，或者允许的j索引逻辑非常复杂，可直接用Numba装饰器加速原生循环，性能和向量化方案基本持平：

import numpy as np
from numba import njit

n = 1000
a = np.random.rand(n)
b = np.random.rand(n)
C = np.zeros((n, n))
allowed_j_per_i = [
    [10, 20] if i % 2 == 0 else [34, 100] 
    for i in range(n)
]

@njit  # 无Python对象模式编译，运行速度接近C
def fill_C(C, a, b, allowed_j_per_i):
    for i in range(len(a)):
        for j in allowed_j_per_i[i]:
            C[i, j] = b[j] - a[i]

fill_C(C, a, b, allowed_j_per_i)

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