CODEWITHSUNDEEP
pythonarraysnumpy
Chris
Chris

Reputation: 13660

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.

Upvotes: 0

Views: 73

Answers (1)

jatha
jatha

Reputation: 934

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

Upvotes: 3

Related Questions