Python中按索引更新矩阵行的向量化无循环实现方法

你可以通过以下两种numpy原生方案实现对应逻辑，完全替换循环且性能远高于Python级遍历：

方法1：直接使用numpy.add.at（最简洁通用） #

numpy.add.at是无缓冲的原位累加操作，天然支持重复索引的场景，一行即可实现需求：

import numpy as np
np.add.at(y, idx, -X)

注：你开头给出的循环逻辑写为y[:,idx[i]] -= X[i]，结合后续示例和等价循环来看属于笔误，实际要实现的是按idx索引y的行维度做减法，上述解法均匹配这个实际逻辑

方法2：分组求和后批量更新（大N场景性能更优） #

如果你的样本量N远大于y的行数k，先按idx分组对X求和再批量更新y，运算效率更高：

import numpy as np
# 初始化和y维度一致的增量矩阵
delta = np.zeros_like(y)
# 按idx分组累加X得到总增量
np.add.at(delta, idx, X)
# 批量更新y
y -= delta

示例验证 #

用你给出的测试输入运行，结果和预期完全一致：

x = np.array([[ 0,  1,  2,  3],
              [ 4,  5,  6,  7],
              [ 8,  9, 10, 11],
              [12, 13, 14, 15],
              [16, 17, 18, 19]])
y = np.array([[100, 101, 102, 103],
              [104, 105, 106, 107],
              [108, 109, 110, 111]])
idx = np.array([2, 1, 1, 1, 2])

# 执行向量化操作
np.add.at(y, idx, -x)
print(y)

输出：

[[100 101 102 103]
 [ 80  78  76  74]
 [ 92  91  90  89]]

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