Python:如何返回连接节点的列表

Question

我有一组节点(存储为nodes = np.array([[x1,y1],[x2,y2],etc]]))和连接节点的元素。每个元素都是两个节点索引的数组。

所以[0,3]是连接节点0和3的元素。

如何返回连接到给定节点索引的所有节点索引的列表？

例如，如果:

element = np.array( [ [3,2] , [1,4] , [1,3] ])
print findConnectedNodes(3) 

应该打印[2,1]

Answer 1

这是一种矢量化方法 -

def findConnectedNodes(a, node_id = 3):
    b = a[(a == node_id).any(1)]
    return b[b!=node_id]

示例运行 -

In [40]: element
Out[40]: 
array([[3, 2],
       [1, 4],
       [1, 3]])

In [41]: findConnectedNodes(element, node_id=3)
Out[41]: array([2, 1])

In [42]: findConnectedNodes(element, node_id=1)
Out[42]: array([4, 3])

In [43]: findConnectedNodes(element, node_id=2)
Out[43]: array([3])

<小时/>

为了提高性能，我们可以使用切片来获取这样一个 2 列成对节点输入数组的 b。此外，我们可以在这里使用更多的 bool 掩码。因此，我们还有两种方法，就像这样 -

def findConnectedNodes_v2(a, node_id = 3):
    b = a[(a[:,0] == node_id) | (a[:,1] == node_id)]
    return b[b!=node_id]

def findConnectedNodes_v3(a, node_id = 3):
    mask2D = a == node_id
    mask1D = mask2D[:,0] | mask2D[:,1]
    return a[mask1D][~mask2D[mask1D]]