Numpy: Clever way of accomplishing v[np.arange(v.shape[0]), col_indices] = 1?

Question

Suppose I have a matrix v:

0 0 0
0 0 0
0 0 0
0 0 0
0 0 0
0 0 0

I also have a vector col_indices = [0,0,1,2,2,1] that indicates which column I should put an 1 for each row of matrix v.

The result of the task, in this case, should be:

1 0 0
1 0 0
0 1 0
0 0 1
0 0 1
0 1 0

The following code works:

v[np.arange(v.shape[0]), col_indices] = 1

But I was wondering if there is a clever way to accomplish this, because in the code above I have to create a vector just to index the matrix and that seems wasteful.

I have also tried the following code, but it doesn't do what I want:

v[:, col_indices] = 1

Answer 1

That IS the cleaver way of doing the indexing.

Look at these timings. Generating that array does not take long. Actually indexing those points takes much longer.

In [208]: x=np.zeros((10000,10000))

In [209]: timeit np.arange(x.shape[0]),np.arange(x.shape[1])
10000 loops, best of 3: 23.5 us per loop

In [210]: timeit x[np.arange(x.shape[0]),np.arange(x.shape[1])]=1
1 loops, best of 3: 1.88 ms per loop

A similar question from a couple of weeks ago is numpy shorthand for taking jagged slice

where the poster wanted something simple like a[:, entries_of_interest] as opposed to a[np.arange(a.shape[0]), entries_of_interest]. I argued that this is just a special case of a[I, J] (for any pair of matching indexing arrays).