python - Numpy 沿轴应用并获取行索引

Question

我有一个二维数组(它实际上非常大并且是另一个数组的 View ):

x = np.array([[0, 1, 2],
          [1, 2, 3],
          [2, 3, 4],
          [3, 4, 5]]
        )

我有一个处理数组每一行的函数:

def some_func(a):
    """
    Some function that does something funky with a row of numbers
    """
    return [a[2], a[0]]  # This is not so funky

np.apply_along_axis(some_func, 1, x)

我正在寻找的是某种调用 np.apply_along_axis 函数的方法，这样我就可以访问行索引(对于正在处理的行)，然后能够处理每一行具有此功能:

def some_func(a, idx):
    """
    I plan to use the index for some logic on which columns to
    return. This is only an example
    """
    return [idx, a[2], a[0]]  # This is not so funky

Answer 1

我有一个 3 维张量的问题，所以我认为发布一个通用的解决方案是值得的，即使用 np.ndenumerate。

    f = lambda indices: #(whatever you'd like to do)

    output = np.empty(M.shape)
    for i, x in np.ndenumerate(M):
        output(i) = f(i)