numpy fancy indexing in multiple dimensions,"print(A.reshape(-1, k)[np.arange(n * m), B.ravel()])"