有没有一种方法可以根据不同索引集的数组，从 2d 数组的对称子矩阵高效构造 3d 数组？

Question

我有一个二维数组

A

和另一个二维数组

I

，其中每一行对应于一组索引，我想用这些索引从

A

中提取对称二维子数组。我想将所有这些子矩阵收集到一个三维数组中，而不需要借助 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

可能是答案，但我尝试过的都不起作用。任何帮助将不胜感激，因为构造这个子矩阵矩阵是我的代码中的一个大瓶颈！

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