нарезка Numpy массив

Question

упражнение в учебнике this говорит:

Сформировать массив 10 х 3 случайных чисел (в диапазоне [0,1]). Для каждой строки выберите номер, ближайший к 0.5.

Используйте abs и argsort, чтобы найти столбец j, ближайший для каждой строки.

Используйте причудливую индексацию для извлечения чисел. (Подсказка:. a[i,j] - массив i должен содержать номера строк, соответствующий материал в j)

Так что я сделал все, но я чувствую, что метод нарезания я использовал (и инициализация b) не вещий на всего:

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

Что другой способ сделать это без использования цикла for?

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

Answer 2

Это технически не избежать for петли, но более вещий способ писать было бы избежать первоначального np.ones назначение и просто использовать список понимание:

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