从两个矩阵中选择行

时间:2017-07-08 09:45:16

标签: python numpy

我正在尝试解决以下问题。我有两个矩阵A和B,我想创建一个新的矩阵C,它由矩阵A和B的行组成,这取决于在数组v中编码的某些条件,即如果v的第i个条目是a那么我希望第i行的第一行是B的第i行,如果它是零那么它应该是第i行。我想出了以下解决方案

C = np.choose(v,A.T,B.T).T

但太慢了。一个明显的坏事是两个转置,但由于np.choose不接受轴参数,我不知道如何摆脱它们。有关快速解决此问题的任何想法吗?

例如,让

A = np.arange(20).reshape([4,5])

B = 10 - A

然后可以想象,人们希望矩阵C是具有最小最大范数的行的矩阵。所以我们让

v = np.sum(A,axis=1)<np.sum(B,axis=1)

然后C是矩阵

C = np.choose(v,[A.T,B.T]).T

array([[10,  9,  8,  7,  6],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19]])

1 个答案:

答案 0 :(得分:1)

似乎是一个很好的设置,可以使用np.where根据掩码/二进制输入数据进行选择操作 -

C = np.where(v[:,None],B,A)

v[:,None]部分基本上将v扩展为可广播的形状AB,以允许broadcasting按照适当的轴进行选择,{{在这种情况下,1}}用于两个axis=0数组。

示例运行 -

2D

如果In [58]: A Out[58]: array([[82, 78, 57], [14, 97, 32], [72, 11, 49], [98, 34, 41], [89, 71, 52], [34, 51, 55], [26, 92, 59]]) In [59]: B Out[59]: array([[55, 67, 50], [49, 64, 21], [34, 18, 72], [24, 61, 65], [56, 59, 23], [44, 77, 13], [56, 55, 58]]) In [62]: v Out[62]: array([1, 0, 0, 0, 0, 1, 1]) In [63]: np.where(v[:,None],B,A) Out[63]: array([[55, 67, 50], [14, 97, 32], [72, 11, 49], [98, 34, 41], [89, 71, 52], [44, 77, 13], [56, 55, 58]]) 严格地仅由v0s组成,请使用1s作为v[:,None]==1的第一个参数。

另一种方法是使用boolean-indexing -

np.where

注意:如果C = A.copy() mask = v==1 C[mask] = B[mask] 已经是布尔数组,请跳过与v的比较以创建蒙版。

运行时测试 -

1