有没有高效方法根据不同索引集从二维数组的对称子矩阵构建三维数组？

Question

我有一个二维数组 A，还有另一个二维数组 I，其中每一行都是一组索引，我想用这些索引从 A 中提取对称的二维子数组。我想把所有这些子矩阵收集到一个三维数组中，而不想使用循环。

比如说，

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 可能是答案，但我尝试的都不管用。如果有人能帮忙就太好了，因为构建这个子矩阵的矩阵是我代码中的一个大瓶颈！

Answer 1

一些索引的小技巧！

B = A[ I[:, :, None], I[:, None, :] ]

验证：

>>> np.array_equal(
...     A[ I[:, :, None], I[:, None, :] ],
...     np.array([A[idx, :][:, idx] for idx in I ])
... )
True