我有一个二维数组A，以及另一个二维数组I，其中每一行对应于一组索引，我想用它们从A中提取对称的二维子array.我希望将所有这些子矩阵收集到一个三维数组中，而不是求助于for循环.
例如:
I = Array([[ 0, 1, 0],
[ 0, 1, 1],
[ 0, 1, 2],
...,
[ 0, 1, 997],
[ 0, 1, 998],
[ 0, 1, 999]], dtype=int32)
因此，A的第一个子矩阵将是对称的3x3子数组A[[0, 1, 0], :][:, [0, 1, 0]].现在，我想定义一个三维数组B，其中有B[0, :, :] = A[I[0, :], :][:, I[0, :]]、B[1, :, :] = A[I[1, :], :][:, I[1, :]]等.
我目前有以下代码
import numpy as np
A = A = np.random.rand(1000,1000)
current_idxs = np.arange(2)
possible_next_idxs = np.arange(1000)
I = np.hstack(( np.tile( current_idxs, (possible_next_idxs.shape[0], 1) ), possible_next_idxs ))
那么一个简单但缓慢的解决方案是
B = np.array([A[idx, :][:, idx] for idx in I ])
这给了我需要的结果，但速度非常慢，必须有一种方法来做我想要的矢量化的方式，但我不确定如何.似乎np.take_along_axis个可能是答案，但我try 过的都没有奏效.任何帮助都将不胜感激，因为构建这个子矩阵矩阵是我的代码中的一大瓶颈！