Как рассчитать сумму всех столбцов двумерного массива (эффективно)

Question

Допустим, у меня есть следующий двумерный массив numpy, состоящий из четырех строк и трех столбцов:

>>> a = numpy.arange(12).reshape(4,3)
>>> print(a)
[[ 0  1  2]
 [ 3  4  5]
 [ 6  7  8]
 [ 9 10 11]]

Какой эффективный способ создать массив 1D, который содержит сумму всех столбцов (например, [18, 22, 26])? Можно ли это сделать без необходимости циклически проходить по всем столбцам?

Answer 1

Проверьте документацию для numpy.sum уделяя особое внимание axis параметр. Для суммирования по столбцам:

>>> import numpy as np
>>> a = np.arange(12).reshape(4,3)
>>> a.sum(axis=0)
array([18, 22, 26])

Или, чтобы суммировать по строкам:

>>> a.sum(axis=1)
array([ 3, 12, 21, 30])

Другие агрегатные функции, такие как numpy.mean, numpy.cumsum а также numpy.std например, также взять axis параметр.

Из предварительного урока Numpy:

Многие унарные операции, такие как вычисление суммы всех элементов в массиве, реализованы как методы ndarray учебный класс. По умолчанию эти операции применяются к массиву, как если бы это был список чисел, независимо от его формы. Однако, указав axis Параметр позволяет применить операцию по указанной оси массива:

Answer 2

Другие варианты суммирования столбцов

numpy.einsum('ij->j', a)

а также

numpy.dot(a.T, numpy.ones(a.shape[0]))

Если количество строк и столбцов одного порядка, все возможности примерно одинаково быстры:

Однако, если есть только несколько столбцов, оба einsum и dot решение значительно превосходит Numpy's sum (обратите внимание на шкалу журнала):

---

Код для воспроизведения сюжетов:

import numpy
import perfplot


def numpy_sum(a):
    return numpy.sum(a, axis=1)


def einsum(a):
    return numpy.einsum('ij->i', a)


def dot_ones(a):
    return numpy.dot(a, numpy.ones(a.shape[1]))


perfplot.show(
    # setup=lambda n: numpy.random.rand(n, n),
    setup=lambda n: numpy.random.rand(n, 3),
    n_range=[2**k for k in range(15)],
    kernels=[numpy_sum, einsum, dot_ones],
    logx=True,
    logy=True,
    xlabel='len(a)',
    )

Answer 3

Использовать axis аргумент:

>> numpy.sum(a, axis=0)
  array([18, 22, 26])

Answer 4

Использование numpy.sum, для вашего случая это

sum = a.sum(axis=0)

Answer 5

Тогда NumPy sum Функция принимает необязательный аргумент оси, который указывает, по какой оси вы хотите получить сумму:

>>> a = numpy.arange(12).reshape(4,3)
>>> a.sum(0)
array([18, 22, 26])

Или, что эквивалентно:

>>> numpy.sum(a, 0)
array([18, 22, 26])

Answer 6

a.sum(0)

должен решить проблему. Это 2d np.array и вы получите сумму всего столбца. axis=0 это измерение, которое указывает вниз и axis=1 тот, который указывает на право.