NumPy - 利用向量化利用另一个数组中的行填充数组的每个下对角线

Question

我有一个形状为

A

的 numpy 数组

(N, M, M)

和另一个形状为

B

的 numpy 数组

(N, M)

。我想用

N

中的行填充子矩阵

A[i, :, :]

的每个

B

对角线。我想使用 Numpy 的矢量化功能尽快完成此操作。当然，可以使用 for 循环，但我想要更高效的东西。

import numpy as np

A = np.random.randn(1000, 50, 50)
B = np.random.randn(1000, 50)

for i in range(1000):
    np.fill_diagonal(A[i, :, :], B[i])

Answer 1

一种可能的方法是生成对角线的索引：

import numpy as np

N, M = 4, 3
A = np.zeros((N, M, M), dtype=int)
B = np.arange(1, N*M+1).reshape((N, M))

idx1, idx2 = np.indices((N, M)).reshape(2, -1)
A[idx1, idx2, idx2] = B.ravel()

输出：

array([[[ 1,  0,  0],
        [ 0,  2,  0],
        [ 0,  0,  3]],

       [[ 4,  0,  0],
        [ 0,  5,  0],
        [ 0,  0,  6]],

       [[ 7,  0,  0],
        [ 0,  8,  0],
        [ 0,  0,  9]],

       [[10,  0,  0],
        [ 0, 11,  0],
        [ 0,  0, 12]]])

索引器：

# idx1
array([0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3])

# idx2
array([0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2])