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.