Методы для работы со всеми столбцами в двумерном массиве.

Question

У меня есть список списков значений (где каждый список значений имеет одинаковую длину), в 2D ndarray. пример

[[0,1,0],
 [0,1,0],
 [0,1,0],
 [1,0,1],
 [1,0,1]]

Я хотел бы выполнить операцию, которая усредняет столбец и возвращает, если среднее значение больше, чем.5. Как бы я достиг этого в NumPy?

Я знаю, что мог бы использовать

for i in range(len(vectors[0])):
    cur_vals = []        
    for j in range(len(vectors)):
        cur_vals.append(vectors[j][i])
    col_operation(cur_vals)

Однако мне сказали, что numpy предоставляет утилиты, подобные R, для этого, поэтому я специально искал это.

Answer 1

Вы можете указать аргумент оси для .mean():

In [180]: a
Out[180]: 
array([[0, 1, 0],
       [0, 1, 0],
       [0, 1, 0],
       [1, 0, 1],
       [1, 0, 1]])

In [194]: a.mean(axis=0)
Out[194]: array([ 0.4,  0.6,  0.4])

Теперь сгенерируйте логическую маску, например:

In [195]: a.mean(axis=0) > .5
Out[195]: array([False,  True, False], dtype=bool)

Чтобы получить столбец со средним значением> .5, используйте логическое индексирование:

In [196]: a.T[a.mean(axis=0) > .5] # you'll need the transpose first
Out[196]: array([[1, 1, 1, 0, 0]])