Numpy 修改 ndarray 对角线

Question

有没有办法在numpy中获取数组的对角线部分的引用？我想把数组的对角线元素除以一个特定的数。

Answer 1

np.fill_diagonal 这个函数运行得非常快：

np.fill_diagonal(a, a.diagonal() / c)

这里的 a 是你的数组，c 是你的因子。在我的电脑上，这种方法的速度和 @kwgoodman 的 a.flat[::n+1] /= c 方法差不多，我觉得这种写法稍微清晰一点（虽然没有那么炫）。

Answer 2

想要快速访问一个方形的 (n,n) numpy 数组的对角线，可以使用 arr.flat[::n+1] 这个方法：

n = 1000
c = 20
a = np.random.rand(n,n)

a[np.diag_indices_from(a)] /= c # 119 microseconds
a.flat[::n+1] /= c # 25.3 microseconds

Answer 3

如果X是你的数组，而c是一个因子，

X[np.diag_indices_from(X)] /= c

可以查看Numpy手册中的diag_indices_from。