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