创建由依赖于索引的二维数组组成的4-D numpy数组

Question

I have a 2-D nx2 array where n is a variable as the code runs. Take n=4 for example, the original array is

[[0, 1],
 [2, 3],
 [4, 5],
 [6, 7]]

I want to create a 4-D nxnx2x2 array where each element of the top level nxn array is an index-dependent 2x2 array made up of the ith and jth row of the original array. For example, at index [0,3] the 2x2 array will be

[[0, 1],
 [6, 7]]

I have tried using the np.fromfunction function and construct the 2x2 array at each coordinate.

new = np.fromfunction(lambda i,j: np.stack((old[i],old[j])), (n,n), dtype=int)

Instead of getting an nxnx2n2 array, I got a 2xnxnxn instead.

Of course, there is the option of using a nested loop and iterating over nxn times but I'd like to do it in a faster way, if possible.

Answer 1

Here is an example by using broadcast_arrays() and stack() :

import numpy as np

a = np.array(
    [[0, 1],
     [2, 3],
     [4, 5],
     [6, 7]])

b = np.stack(np.broadcast_arrays(a[:, None, :], a[None, :, :]), axis=-2)
print(b[0, 3])