Numpy: set one specific element of each column based on indexing by array

Question

On a smaller scale compared to what I need, here's an example of what I'm looking to do:

>>> a
array([[  21,   22,   23,   24,   25,   26,   27],
       [  56,   57,   58,   59,   60,   61,   62],
       [  14,   15,   16,   17,   18,   19,   20],
       [   7,    8,    9, 1010,   11,   12,   13],
       [  42,   43,   44,   45,   46,   47,   48],
       [  63,   64,   65,   66,   67,   68,   69],
       [   0,    1,    2,    3,    4,    5,    6],
       [  49,   50,   51,   52,   53,   54,   55],
       [  28,   29,   30,   31,   32,   33,   34],
       [  35,   36,   37,   38,   39,   40,   41]])
>>> indices = a.argmax(axis=0)
>>> indices
array([5, 5, 5, 3, 5, 5, 5])
>>> b = np.zeros(a.shape)
>>> b[indices] = 1.0
>>> b # below is the actual output, not what I want
array([[ 0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 1.,  1.,  1.,  1.,  1.,  1.,  1.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 1.,  1.,  1.,  1.,  1.,  1.,  1.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.]])

But what I actually need is:

>>> b
array([[ 0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  1.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 1.,  1.,  1.,  0.,  1.,  1.,  1.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.]])

Numpy indexing can get extremely complicated and it's a little difficult to put the above into words, so hopefully someone can understand what I'm looking for. Essentially it's to set a 1 wherever there's a max of a column and zero elsewhere. How would I go about doing this?

Answer 1

From the docs:

If the number of objects in the selection tuple is less than N , then : is assumed for any subsequent dimensions.

In your selection there only one array, so you get every row from indices to be equal to 1. To overcome that, you need column indices. I guess this will do the trick:

b[indices, np.arange(a.shape[1])] = 1.0