给定两个 numpy 向量a
和b
,我想创建一个 2D 数组A
,如果a[i]==b[j]
,A[i,j]=True
。我可以执行以下 for 循环。但我想知道我是否可以在没有 for 循环的情况下做到这一点以更快地获得结果。
def f(x, y):
d = np.zeros((len(x),len(y)))
for i in range(len(x)):
d[i, np.argwhere(y==x[i])]=1
return d
a = np.array([1,2,3,1,3,4])
b = np.array([1,1,2,2,2,3,4,5])
f(a,b)
是:
array([[1., 1., 0., 0., 0., 0., 0., 0.],
[0., 0., 1., 1., 1., 0., 0., 0.],
[0., 0., 0., 0., 0., 1., 0., 0.],
[1., 1., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 1., 0., 0.],
[0., 0., 0., 0., 0., 0., 1., 0.]])
In [296]: a = np.array([1,2,3,1,3,4])
...: b = np.array([1,1,2,2,2,3,4,5])
In [297]: a[:,None]==b
Out[297]:
array([[ True, True, False, False, False, False, False, False],
[False, False, True, True, True, False, False, False],
[False, False, False, False, False, True, False, False],
[ True, True, False, False, False, False, False, False],
[False, False, False, False, False, True, False, False],
[False, False, False, False, False, False, True, False]])
In [298]: _.astype(int)
Out[298]:
array([[1, 1, 0, 0, 0, 0, 0, 0],
[0, 0, 1, 1, 1, 0, 0, 0],
[0, 0, 0, 0, 0, 1, 0, 0],
[1, 1, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 1, 0, 0],
[0, 0, 0, 0, 0, 0, 1, 0]])
或带where
:
In [301]: d=np.zeros((len(a),len(b)),int)
In [302]: np.where(a[:,None]==b)
Out[302]: (array([0, 0, 1, 1, 1, 2, 3, 3, 4, 5]), array([0, 1, 2, 3, 4, 5, 0, 1, 5, 6]))
In [303]: d[_]=1
In [304]: d
Out[304]:
array([[1, 1, 0, 0, 0, 0, 0, 0],
[0, 0, 1, 1, 1, 0, 0, 0],
[0, 0, 0, 0, 0, 1, 0, 0],
[1, 1, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 1, 0, 0],
[0, 0, 0, 0, 0, 0, 1, 0]])