在 NumPy 数组中推广切片操作

这个问题基于这个较旧的问题:

This question is based on this older question:

给定一个数组:

In [122]: arr = np.array([[1, 3, 7], [4, 9, 8]]); arr
Out[122]: 
array([[1, 3, 7],
       [4, 9, 8]])

并给出它的索引:

In [127]: np.indices(arr.shape)
Out[127]: 
array([[[0, 0, 0],
        [1, 1, 1]],

       [[0, 1, 2],
        [0, 1, 2]]])

我如何才能将它们整齐地堆叠在一起形成一个新的二维数组?这就是我想要的:

How would I be able to stack them neatly one against the other to form a new 2D array? This is what I'd like:

array([[0, 0, 1],
       [0, 1, 3],
       [0, 2, 7],
       [1, 0, 4],
       [1, 1, 9],
       [1, 2, 8]])

Divakar 的这个解决方案是我目前用于二维数组的解决方案:

This solution by Divakar is what I currently use for 2D arrays:

def indices_merged_arr(arr):
    m,n = arr.shape
    I,J = np.ogrid[:m,:n]
    out = np.empty((m,n,3), dtype=arr.dtype)
    out[...,0] = I
    out[...,1] = J
    out[...,2] = arr
    out.shape = (-1,3)
    return out

现在，如果我想传递一个 3D 数组，我需要修改这个函数:

Now, if I wanted to pass a 3D array, I need to modify this function:

def indices_merged_arr(arr):
    m,n,k = arr.shape   # here
    I,J,K = np.ogrid[:m,:n,:k]   # here
    out = np.empty((m,n,k,4), dtype=arr.dtype)   # here
    out[...,0] = I
    out[...,1] = J
    out[...,2] = K     # here
    out[...,3] = arr
    out.shape = (-1,4)   # here
    return out

但此函数现在仅适用于 3D 数组 - 我无法将 2D 数组传递给它.

But this function now works for 3D arrays only - I can't pass a 2D array to it.

是否有某种方法可以将其概括为适用于任何维度?这是我的尝试:

Is there some sort of way I can generalise this to work for any dimension? Here's my attempt:

def indices_merged_arr_general(arr):
    tup = arr.shape   
    idx = np.ogrid[????]   # not sure what to do here....
    out = np.empty(tup + (len(tup) + 1, ), dtype=arr.dtype) 
    for i, j in enumerate(idx):
        out[...,i] = j
    out[...,len(tup) - 1] = arr
    out.shape = (-1, len(tup)
    return out

我遇到了这条线的问题:

I'm having trouble with this line:

idx = np.ogrid[????]   

我怎样才能让它工作?