如何在需要唯一性和非唯一性检查的Python循环中向量化这个循环？

Question

我有一个大小为3xn的numpy数组（edges），它正好包含两行，其中行的前两个元素匹配，第三个元素可能匹配，也可能不匹配。我正在根据这个标准计算一些东西。我基于for循环的代码是这样的

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]]

Answer 1

我还没有完全了解您是如何设置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]])