I am trying to obtain the value corresponding to column b[i] for each row i in A
Can I do this without using the for loop?
A = np.array([[35, 2, 23, 22], [44, 21, 15, 4], [44, 21, 15, 4], [37, 4, 17, 41], [33, 4, 4, 18], [35, 2, 23, 22]])
b = np.array([0,1,1,2,3,0])
C = zeros(len(b),1)
for i in range(6):
C[i] = A[i][b[i]]
Since you want to sequentially index the rows of A, you can index with np.arange(len(A)) in addition to b to get your desired output:
A[np.arange(len(A)), b]
# array([35, 21, 21, 17, 18, 35])
Showing how this works:
# A np.arange(len(A)) b
array([[35, 2, 23, 22], [0, 0]
[44, 21, 15, 4], [1, 1]
[44, 21, 15, 4], [2, 1]
[37, 4, 17, 41], [3, 2]
[33, 4, 4, 18], [4, 3]
[35, 2, 23, 22]]) [5, 0]
How about this?
A[:,b].diagonal()
Out[133]: array([35, 21, 21, 17, 18, 35])
