Chris Chris - 5 months ago 20
Python Question

Index with an unknown number of indices

I have one numpy array with a N dimensions where N can vary:

arr1 = np.array([0,0,1,0])
arr2 = np.array([[0,0],[0,0],[0,1],[0,0]])


I need to be able to pass an array and tuple/int and get the integer in that location returned:

assert myfunc(arr1, 2) == 1
assert myfunc(arr2, (2, 1)) == 1


I feel like this has an obvious and simple answer that is just not clicking for some reason.

Answer

I'm not sure I understood your question correctly; but, you can directly use tuple as an index to retrieve the specific elements.

For instance;

arr1 = np.array([0,0,1,0])
arr2 = np.array([[0,0],[0,0],[0,1],[0,0]])
arr3 = np.array([[[0,0],[0,0],[0,1],[0,0]], [[0,0],[0,0],[0,1],[0,0]]])

assert arr1[2] == 1
assert arr2[(2, 1)] == 1
assert arr3[(1,2,1)] == 0
Comments