我有一个大小为edges
的numpy数组(3xn
),它包含两行,其中行的前两个元素匹配,第三个元素可能匹配,也可能不匹配。我正在根据这个标准计算一些东西。我的基于循环的代码沿着这些行开始
mat = np.zeros((edges.shape[0],2),dtype=np.int64)
counter = 0;
for i in range(edges.shape[0]):
for j in range(i+1,edges.shape[0]):
if edges[i,0]==edges[j,0] and edges[i,1]==edges[j,1] and edges[i,2] != edges[j,2]:
mat[2*counter,0] = i % nof
mat[2*counter,1] = i // nof
mat[2*counter+1,0] = j % nof
mat[2*counter+1,1] = j // nof
counter +=1
break
其中nof
是特定数字。如何使用numpy加速此代码?我无法使用np.unique
,因为此代码需要唯一性以及非唯一性检查。
例如,给定:
edges = np.array([
[1,2,13],
[4,5,15],
[5,6,18],
[1,2,12],
[4,5,15],
[5,6,18],
])
其中每行的前两个元素可以在另一行中找到(也就是说它们重复两次)和nof=1
,上面的代码给出了以下结果
[[0 0]
[0 3]
[0 0]
[0 0]]
答案 0 :(得分:2)
我还没有意识到你如何设置mat
,但我怀疑前两列的lexsorting可能有所帮助:
In [512]: edges = np.array([
...: [1,2,13],
...: [4,5,15],
...: [5,6,18],
...: [1,2,12],
...: [4,5,15],
...: [5,6,18],
...: ])
...:
In [513]: np.lexsort((edges[:,1],edges[:,0]))
Out[513]: array([0, 3, 1, 4, 2, 5], dtype=int32)
In [514]: edges[_,:] # sedges (below)
Out[514]:
array([[ 1, 2, 13],
[ 1, 2, 12],
[ 4, 5, 15],
[ 4, 5, 15],
[ 5, 6, 18],
[ 5, 6, 18]])
现在所有匹配的行都在一起。
如果总共有2个匹配,则可以收集这些对并将其重新整形为2列数组。
In [516]: sedges[:,2].reshape(-1,2)
Out[516]:
array([[13, 12],
[15, 15],
[18, 18]])
或者你仍然可以迭代,但你不必检查尽可能远。
排序列表上的 argsort
返回反向排序:
In [519]: np.lexsort((edges[:,1],edges[:,0]))
Out[519]: array([0, 3, 1, 4, 2, 5], dtype=int32)
In [520]: np.argsort(_)
Out[520]: array([0, 2, 4, 1, 3, 5], dtype=int32)
In [521]: sedges[_,:]
Out[521]:
array([[ 1, 2, 13],
[ 4, 5, 15],
[ 5, 6, 18],
[ 1, 2, 12],
[ 4, 5, 15],
[ 5, 6, 18]])