Multi-dimensional array indexing using a single dimensional array in Python

Question

I have two dimensional array, X of size (500,10) and a single dimensional index array Y whose size is 500 whose each entry is an index of correct value column of corresponding row of X, e.g, y(0) is 2 then it means column 2 of first row of X is correct, similarly y(3) = 4 means Row 3 and Col 4 of X has correct value.

I want to get all the correct values from X using index array Y without using any loops, i.e, using vectorization and in this case the output should be (500,1). But when i do X[:,y] then it gives output (500,500). Can someone help me how to correctly index array X using Y, plz.

Thank you all for the help.

Answer 1

Another option is multidimensional list-of-locations indexing:

import numpy as np

ncol = 10  # 10 in your case
nrow = 500  # 500 in your case
# just creating some test data:
x = np.arange(ncol*nrow).reshape(nrow,ncol)
y = (ncol * np.random.random_sample((nrow, 1))).astype(int)

print(x)
print(y)
print(x[np.arange(nrow),y.T].T)

The syntax is explained here. You basically need an array of indices for each dimension. In the first dimension this is simply [0,...,500] in your case and the second dimension is your y-array. We need to transpose it (.T), because it has to have the same shape as the first and the output array. The second transposition is not really needed, but gives you the shape you want.

EDIT:

The question of performance came up and I tried the three methods mentioned so far. You'll need line_profiler to run the following with

kernprof -l -v tmp.py

where tmp.py is:

import numpy as np

@profile
def calc(x,y):
    z = np.arange(nrow)
    a = x[z,y.T].T  # mine, with the suggested speed up
    b = x[:,y].diagonal().T  # Christoph Terasa
    c = np.array([i[j] for i, j in zip(x, y)])  # tobias_k

    return (a,b,c)

ncol = 5  # 10 in your case
nrow = 10  # 500 in your case

x = np.arange(ncol*nrow).reshape(nrow,ncol)
y = (ncol * np.random.random_sample((nrow, 1))).astype(int)

a, b, c = calc(x,y)
print(a==b)
print(b==c)

The output for my python 2.7.6:

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    3                                           @profile
    4                                           def calc(x,y):
    5         1            4      4.0      0.1      z = np.arange(nrow)
    6         1           35     35.0      0.8      a = x[z,y.T].T
    7         1         3409   3409.0     76.7      b = x[:,y].diagonal().T
    8       501          995      2.0     22.4      c = np.array([i[j] for i, j in zip(x, y)])
    9                                           
    10         1            1      1.0      0.0      return (a,b,c)

Where %Time or Time are the relevant columns. I don't know how to profile memory consumption, someone else would have to do that. For now it looks like my solution is the fastest for the requested dimensions.

Answer 2

While not really intuitive from a syntactic perspective

X[:,Y].diagonal()[0]

will give you the values you're looking for. The fancy indexing selects from each row of X all values in Y, and diagonal selects only those at the indexes where i == j. The indexing with [0] at the end just flattens the 2d array.