Fancy indexing calculation of adjacency matrix from adjacency list

Question

Problem:

I want to calculate at several times the adjacency matrix A_ij given the adjacency list E_ij, where E_ij[t,i] = j gives the edge from i to j at time t.

I can do it with the following code:

import numpy as np

nTimes = 100
nParticles = 10
A_ij = np.full((nTimes, nParticles, nParticles), False)
E_ij = np.random.randint(0, 9, (100, 10))

for t in range(nTimes):
    for i in range(nParticles):
        A_ij[t, i, E_ij[t,i]] = True

Question:

How can I do it in a vectorized way, either with fancy indexing or using numpy functions such as np.take_along_axis?

---

What I tried:

I expected this to work:

A_ij[:,np.arange(nParticles)[None,:,None], E_ij[:,None,np.arange(nParticles)]] = True

But it does not.

---

Related to: Trying to convert adjacency list to adjacency matrix in Python

Answer 1

I think this might work:

import numpy as np

nTimes = 100
nParticles = 10
A_ij = np.full((nTimes, nParticles, nParticles), False)
E_ij = np.random.randint(0, 9, (100, 10))

np.put_along_axis(A_ij, E_ij[..., None], True, axis=2)