Accessing indices from functions acting on entire arrays-CodePudding

I want a function, f, that given an (n, 1) numpy array returns an (n, 1) numpy array where the values are those of the original plus their 0-axis index.

Ex:

f(np.array([[0],[0]])) = np.array([[0],[1]])

f(np.array([[4],[2],[3]])) = np.array([[4],[3],[5]])

Is there some command, cmd(x), that returns the index of the value the function is operating on? The code for f would look like:

def f(x):
    return x   cmd(x)

EDIT: Fixed second example

CodePudding user response：

If you want something that'll work no matter what the shape of x, you want something like:

np.indices(x.shape)[0] 1

This can be optimized to build a smaller array if you're using a large array.