Question

The exercise in this tutorial says:

Generate a 10 x 3 array of random numbers (in range [0,1]). For each row, pick the number closest to 0.5.

Use abs and argsort to find the column j closest for each row.

Use fancy indexing to extract the numbers. (Hint: a[i,j] – the array i must contain the row numbers corresponding to stuff in j.)

So I did everything, but I feel the slicing method I used (and the initialisation of b) is not pythonic at all:

a = np.random.rand(10,3)

mask = np.argmin(abs(a-0.5), axis = 1)

b = np.ones(mask.size)

for j in range(0,mask.size):
    b[j] = a[j,mask[j]]

What is the other way of doing this without using the for loop?

Answer 1

import numpy as np
a = np.random.rand(10,3)
b = np.argmin(abs(a - .5), axis=1).choose(a.T)

# a
array([[ 0.97272372,  0.45351387,  0.19105835],
       [ 0.27895897,  0.12438789,  0.64857335],
       [ 0.05298066,  0.58122882,  0.805319  ],
       [ 0.39952727,  0.77728036,  0.65742471],
       [ 0.36522802,  0.06938552,  0.6595684 ],
       [ 0.9030323 ,  0.08965774,  0.01823633],
       [ 0.30996923,  0.53400339,  0.87600912],
       [ 0.17953532,  0.4888832 ,  0.0746074 ],
       [ 0.09052476,  0.47397504,  0.30317449],
       [ 0.31851577,  0.68135476,  0.38335483]])

# b
array([ 0.45351387,  0.64857335,  0.58122882,  0.39952727,  0.36522802,
        0.9030323 ,  0.53400339,  0.4888832 ,  0.47397504,  0.38335483])