Filter out index from 2-D numpy array
Question:
Say I have a 2-D numpy array,
A = [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]] # Shape: (3, 4)
and I have another 1-D array,
B = [0, 2, 1] # Shape: (3,)
I want to extract an element from the index of A’s ith row corresponding to B’s ith element.
Meaning, if we consider the 3nd element in B, output will be 2nd element of 3rd row of A (that is, 9).
Thus, the desired output is:
C = [0, 6, 9]
Now, this can easily be done using a for loop, however, I am looking for some other optimized ways of doing this.
I tried np.take(), however it’s not working as desired.
Answers:
If your B array is the position along the second axis which you’d like for each element in the first axis, just provide a corresponding set of indices into the first dimension:
In [4]: A[range(A.shape[0]), B]
Out[4]: array([0, 6, 9])
This is equivalent to:
In [5]: A[[0, 1, 2], [1, 2, 0]]
Out[5]: array([0, 6, 9])
Simple enough with advanced indexing
C = A[np.arange(B.size), B]
which, for your example, yields [0, 6, 9] which is the correct result as @MichaelDelgado points out.
Answer: Use List Comprehension with Enumerate
You could list comprehension, with enumerate which will give you a list of touples in the syntax of (index, value). So for instance, enumerate([0,2]) will result in [(0,0), (1,2)].
So, in order to do so, it would be
In: C = [A[point[0]][point[1]] for point in enumerate(B)]
Out: [0,6,9]
which would be the answer you’re looking for. I believe the 8 you have in your supposed answer is incorrect.
Say I have a 2-D numpy array,
A = [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]] # Shape: (3, 4)
and I have another 1-D array,
B = [0, 2, 1] # Shape: (3,)
I want to extract an element from the index of A’s ith row corresponding to B’s ith element.
Meaning, if we consider the 3nd element in B, output will be 2nd element of 3rd row of A (that is, 9).
Thus, the desired output is:
C = [0, 6, 9]
Now, this can easily be done using a for loop, however, I am looking for some other optimized ways of doing this.
I tried np.take(), however it’s not working as desired.
If your B array is the position along the second axis which you’d like for each element in the first axis, just provide a corresponding set of indices into the first dimension:
In [4]: A[range(A.shape[0]), B]
Out[4]: array([0, 6, 9])
This is equivalent to:
In [5]: A[[0, 1, 2], [1, 2, 0]]
Out[5]: array([0, 6, 9])
Simple enough with advanced indexing
C = A[np.arange(B.size), B]
which, for your example, yields [0, 6, 9] which is the correct result as @MichaelDelgado points out.
Answer: Use List Comprehension with Enumerate
You could list comprehension, with enumerate which will give you a list of touples in the syntax of (index, value). So for instance, enumerate([0,2]) will result in [(0,0), (1,2)].
So, in order to do so, it would be
In: C = [A[point[0]][point[1]] for point in enumerate(B)]
Out: [0,6,9]
which would be the answer you’re looking for. I believe the 8 you have in your supposed answer is incorrect.