python - 给定矩阵列表，对它们求和但忽略某些矩阵的行

Question

举个例子

A = np.array([[a1, a2],
              [a3, a4]])
B = np.array([[b1, b2],
              [b3, b4]])
C = np.array([[c1, c2],
              [c3, c4]]),

让 l = [A, B, C]，让

I = np.array([[1, 0, 1],
              [0, 1, 1]]).

我要计算矩阵R

R = np.array([[a1 + c1, a2 + c2],
              [b3 + c3, b4 + c4]])

---

即一般来说，我有一个 k nxm 矩阵列表存储在 l 和一个索引矩阵 I 维度 nxk 指定结果 R 的每一行 l 中 k 矩阵中的哪一个> 应该用于计算 R 的那一行。上面，I 的第一行是 [1, 0, 1] 因此使用了 A 和 C在计算第一行的R。

我可以为此使用哪个 numpy 函数？

Answer 1

假设 l 作为 (n,m) 形数组的输入列表，您可以使用非常高效的 np.einsum , 像这样 -

np.einsum('ij,jik->ik',I,l)

sample 运行-

In [75]: # Inputs
    ...: a1,a2,a3,a4 = 4,6,8,2
    ...: b1,b2,b3,b4 = 5,11,4,3
    ...: c1,c2,c3,c4 = 6,99,2,5
    ...: 
    ...: A = np.array([[a1, a2],
    ...:               [a3, a4]])
    ...: B = np.array([[b1, b2],
    ...:               [b3, b4]])
    ...: C = np.array([[c1, c2],
    ...:               [c3, c4]])
    ...: l = [A,B,C]
    ...: 
    ...: I = np.array([[1, 0, 1],
    ...:               [0, 1, 1]])
    ...: # Expected output
    ...: R = np.array([[a1 + c1, a2 + c2],
    ...:               [b3 + c3, b4 + c4]])
    ...: 

In [76]: R
Out[76]: 
array([[ 10, 105],
       [  6,   8]])

In [77]: np.einsum('ij,jik->ik',I,l)
Out[77]: 
array([[ 10, 105],
       [  6,   8]])