Index numpy arrays columns by another numpy array

Question

I am trying to index a 2d matrix in numpy so that I can get all rows but only particular columns given by another numpy array. It's something as following:

a = [0,1,1,2,0,2,1]

d = [[1,2,3],[1,2,3],[1,2,3],[1,2,3],[1,2,3],[1,2,3],[1,2,3]]

I want to get all rows from d such that column is given by a. So for above example I want,

t = [1,2,2,3,1,3,2]

I tried some of the methods given on numpy documentation but am not able to get it.

I think this is doable in matlab without any iteration. Can I do this is python without looping over something?

Answer 1

This can be done with advanced indexing:

>>> a = numpy.array([0, 1, 1, 2, 0, 2, 1])
>>> d = numpy.array([[1,2,3],[1,2,3],[1,2,3],[1,2,3],[1,2,3],[1,2,3],[1,2,3]])
>>> d[numpy.arange(d.shape[0]), a]
array([1, 2, 2, 3, 1, 3, 2])

For arrays a, b, and c where b and c have integer dtype and b.shape == c.shape, advanced indexing d = a[b, c] gives d[i] == a[b[i], c[i]].