Indexing one array by another in numpy,"A[np.arange(A.shape[0])[:, (None)], B]"