fgregg fgregg - 6 months ago 50
Python Question

Accumulate values of "neigborhood" from edgelist with numpy

I have a undirected network where each node can be one of k types. For each node i, I need to calculate the number of neighbors that node i has of each type.

Right now I am representing the edges with an edgelist where the columns are indexes of the nodes. The nodes are represented as a n x k matrix, where each column represents a node type. If a node is of type k then the kth column's value is 1, 0 otherwise.

Here's my current code, which is correct, but too slow.

# example nodes and edges, both typically much longer
nodes = np.array([[0, 0, 1],
[0, 1, 0],
[1, 0, 0]])
edges = np.array([[0, 1],
[1, 2]])

neighbors = np.zeros_like(nodes)

for i, j in edges:
neighbors[i] += nodes[j]
neighbors[j] += nodes[i]

Is there some clever numpy that would allow me to avoid this for loop? If the best way to do this is with an adjacency matrix, that is also acceptable.


You can simply use np.add.at -

out = np.zeros_like(nodes)
np.add.at(out, edges[:,0],nodes[edges[:,1]])
np.add.at(out, edges[:,1],nodes[edges[:,0]])