如何用fancy indexing从邻接列表向量化计算邻接矩阵？

给定邻接列表E_ij（其中E_ij[t,i] = j表示时刻t下从节点i到节点j的边），需要在多个时刻计算邻接矩阵A_ij。

目前已通过以下嵌套循环代码实现该功能：

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

如何采用向量化方式实现该功能？例如使用fancy indexing或numpy的np.take_along_axis函数。

本人尝试了以下代码，但未能成功：

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

方法1：Fancy Indexing广播赋值 #

通过构造时刻和节点的索引数组，利用numpy的广播特性直接定位赋值：

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

# 生成时刻索引，形状(100,1)，用于和节点索引广播
t_idx = np.arange(nTimes)[:, np.newaxis]
# 生成节点i的索引，形状(1,10)，用于和时刻索引广播
i_idx = np.arange(nParticles)[np.newaxis, :]

# 三个索引数组广播后均为(100,10)，精准定位每个t,i对应的j位置
A_ij[t_idx, i_idx, E_ij] = True

方法2：使用np.put_along_axis函数 #

利用put_along_axis在指定轴上赋值，无需手动构造广播索引：

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

# 将E_ij扩展为(100,10,1)，匹配A_ij的第三维
target_pos = E_ij[:, :, np.newaxis]
# 创建同形状的True掩码，用于赋值
mask = np.full_like(target_pos, True)

# 在第三维（节点j的维度）上，把mask的值放到target_pos指定的位置
np.put_along_axis(A_ij, target_pos, mask, axis=2)

内容的提问来源于stack exchange，提问作者Puco4