Сумма массива по числу в numpy

Предполагая, что у меня есть массив numpy, например: [1,2,3,4,5,6] и другой массив: [0,0,1,2,2,1] Я хочу суммировать элементы в первом массиве по группе (второй массив) и получить результаты n-групп в порядке порядкового номера группы (в этом случае результат будет [3, 9, 9]). Как это сделать в numpy?

Ответы

Ответ 1

Здесь существует более одного способа сделать это, но здесь одним из способов:

import numpy as np
data = np.arange(1, 7)
groups = np.array([0,0,1,2,2,1])

unique_groups = np.unique(groups)
sums = []
for group in unique_groups:
    sums.append(data[groups == group].sum())

Вы можете векторизовать вещи так, чтобы там вообще не было цикла, но я бы рекомендовал против него. Он становится нечитаемым и потребует пару временных временных массивов 2D, которые могут потребовать больших объемов памяти, если у вас много данных.

Изменить: Здесь вы можете полностью векторизовать. Имейте в виду, что это может (и, скорее всего, будет) медленнее, чем версия выше. (И может быть лучший способ его векторизации, но это поздно, и я устал, так что это просто первое, что появляется в моей голове...)

Однако имейте в виду, что это плохой пример... Вам действительно лучше (как с точки зрения скорости, так и с точки зрения удобочитаемости) с помощью цикла выше...

import numpy as np
data = np.arange(1, 7)
groups = np.array([0,0,1,2,2,1])

unique_groups = np.unique(groups)

# Forgive the bad naming here...
# I can't think of more descriptive variable names at the moment...
x, y = np.meshgrid(groups, unique_groups)
data_stack = np.tile(data, (unique_groups.size, 1))

data_in_group = np.zeros_like(data_stack)
data_in_group[x==y] = data_stack[x==y]

sums = data_in_group.sum(axis=1)

Ответ 2

Это векторный метод выполнения этой суммы, основанный на реализации numpy.unique. По моим таймингам это до 500 раз быстрее, чем метод цикла и в 100 раз быстрее, чем метод гистограммы.

def sum_by_group(values, groups):
    order = np.argsort(groups)
    groups = groups[order]
    values = values[order]
    values.cumsum(out=values)
    index = np.ones(len(groups), 'bool')
    index[:-1] = groups[1:] != groups[:-1]
    values = values[index]
    groups = groups[index]
    values[1:] = values[1:] - values[:-1]
    return values, groups

Ответ 3

Функция numpy bincount была сделана именно для этой цели, и я уверен, что она будет намного быстрее, чем другие методы для всех размеров входов:

data = [1,2,3,4,5,6]
ids  = [0,0,1,2,2,1]

np.bincount(ids, weights=data) #returns [3,9,9] as a float64 array

i-й элемент вывода представляет собой сумму всех элементов data, соответствующих "id" i.

Надеюсь, что это поможет.

Ответ 4

Если группы индексируются последовательными целыми числами, вы можете злоупотреблять функцией numpy.histogram(), чтобы получить результат:

data = numpy.arange(1, 7)
groups = numpy.array([0,0,1,2,2,1])
sums = numpy.histogram(groups, 
                       bins=numpy.arange(groups.min(), groups.max()+2), 
                       weights=data)[0]
# array([3, 9, 9])

Это позволит избежать любых циклов Python.

Ответ 5

Я пробовал скрипты от всех, и мои соображения:

Joe: Будет работать только в том случае, если у вас мало групп.

kevpie: слишком медленно из-за циклов (это не пифонический путь)

Bi_Rico и Sven: выполнить хорошо, но будет работать только для Int32 (если сумма перейдет через 2 ^ 32/2, она не удастся)

Алекс: самый быстрый, полезный для суммы.

Но если вы хотите больше гибкости и возможности группировать по другой статистике, используйте SciPy:

from scipy import ndimage

data = np.arange(10000000)
groups = np.arange(1000).repeat(10000)
ndimage.sum(data, groups, range(1000))

Это хорошо, потому что у вас много статистики для группировки (сумма, значение, дисперсия,...).

Ответ 6

Ты ошибаешься! Лучший способ сделать это:

a = [1,2,3,4,5,6]
ix = [0,0,1,2,2,1]
accum = np.zeros(np.max(ix)+1)
np.add.at(accum, ix, a)
print accum
> array([ 3.,  9.,  9.])

Ответ 7

Реализация чистого python:

l = [1,2,3,4,5,6]
g = [0,0,1,2,2,1]

from itertools import izip
from operator import itemgetter
from collections import defaultdict

def group_sum(l, g):
    groups = defaultdict(int)
    for li, gi in izip(l, g):
        groups[gi] += li
    return map(itemgetter(1), sorted(groups.iteritems()))

print group_sum(l, g)

[3, 9, 9]

Ответ 8

Я заметил тег numpy, но если вы не против использования pandas, эта задача станет однострочным:

import pandas as pd
import numpy as np

data = np.arange(1, 7)
groups = np.array([0, 0, 1, 2, 2, 1])

df = pd.DataFrame({'data': data, 'groups': groups})

Итак, df выглядит следующим образом:

   data  groups
0     1       0
1     2       0
2     3       1
3     4       2
4     5       2
5     6       1

Теперь вы можете использовать функции groupby() и sum()

print df.groupby(['groups'], sort=False).sum()

который дает желаемый результат

        data
groups      
0          3
1          9
2          9

По умолчанию кадр данных будет отсортирован, поэтому я использую флаг sort=False, который может улучшить скорость для огромных фреймов данных.