python - 从邻接表计算邻接矩阵的花式索引

Question

问题:

我想在给定邻接列表E_i​​j的情况下多次计算邻接矩阵A_ij，其中E_i​​j[t,i] = j在时间 t 给出从 i 到 j 的边。

我可以用下面的代码来完成:

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

问题:

我怎样才能以矢量化的方式做到这一点，既可以使用精美的索引，也可以使用诸如 np.take_along_axis 之类的 numpy 函数？

---

我尝试了什么:

我希望它能工作:

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

但事实并非如此。

---

相关:Trying to convert adjacency list to adjacency matrix in Python

Answer 1

我认为这可能有效:

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)