Slicing a numpy array

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?

eumiro · Accepted Answer · 2012-06-20 10:58:24Z

5

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])

answered Jun 20, 2012 at 10:58

eumiro

214k36 gold badges307 silver badges264 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

DSM Over a year ago

Clever! I would have used b = a[range(len(mask)), mask] instead (which I think is what the hint was suggesting) but this is neat.

abdulhaq-e Over a year ago

Python never ceases to amaze me :D. This is for people like me who need to see other examples on how choose works. Here

mVChr · Accepted Answer · 2012-06-20 10:47:46Z

0

It doesn't technically avoid the for loop, but a more pythonic way to write it would be to avoid the initial np.ones assignment and just use a list comprehension:

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

answered Jun 20, 2012 at 10:47

mVChr

50.3k11 gold badges111 silver badges105 bronze badges

Collectives™ on Stack Overflow

Slicing a numpy array

2 Answers 2

2 Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

2 Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related