Средние строки с повторяющимся первым элементом в двумерном массиве

Question

У меня есть такой массив:

[[1, 2, 3, 4],
 [2, 2, 3, 4],
 [1, 4, 5, 6],
 [3, 2, 3, 4]]

Что я хочу сделать, так это усреднить строки каждой повторяющейся первой записи. В этом примере усреднение второго до последнего столбца первой и третьей строк. Правильный вывод должен быть:

[[1, 3, 4, 5],
 [2, 2, 3, 4],
 [3, 2, 3, 4]]

Этот ответ Среднее количество записей с повторяющимся первым элементом в двумерном массиве numpy очень сильно связано с тем, что я хочу сделать.

Answer 1

Вот один из способов:

l = [np.argwhere(i==a[:,0]).ravel() for i in np.unique(a[:,0])]
np.array([a[i].mean(axis=0).tolist() for i in l])

array([[1., 3., 4., 5.],
       [2., 2., 3., 4.],
       [3., 2., 3., 4.]])

Объяснение

Что касается понимания первого списка, np.arwhere будет возвращать массив с индексами строк с общим первым элементом для каждого unique элемента

l = [np.argwhere(i==a[:,0]).ravel() for i in np.unique(a[:,0])]
[array([0, 2]), array([1]), array([3])]

Теперь вам просто нужно нарезать a по каждому массиву в l и усреднить его с помощью:

np.array([a[i].mean(axis=0).tolist() for i in l])

array([[1., 3., 4., 5.],
       [2., 2., 3., 4.],
       [3., 2., 3., 4.]])

Answer 2

Подход № 1

Вот векторизованный с NumPy -

def groupby_mean(a):
    # Sort array by groupby column
    b = a[a[:,0].argsort()]

    # Get interval indices for the sorted groupby col
    idx = np.flatnonzero(np.r_[True,b[:-1,0]!=b[1:,0],True])

    # Get counts of each group and sum rows based on the groupings & hence averages
    counts = np.diff(idx)
    avg = np.add.reduceat(b[:,1:],idx[:-1],axis=0)/counts.astype(float)[:,None]

    # Finally concatenate for the output in desired format
    return np.c_[b[idx[:-1],0],avg]

Подход № 2

Еще один векторизованный, использующий matrix-multiplication -

def groupby_mean_matmul(a):
    unq = np.unique(a[:,0])
    m = a[:,0,None] == unq
    return np.c_[unq, m.T.dot(a[:,1:])/m.sum(0)[:,None].astype(float)]

Пробный прогон -

In [51]: a
Out[51]: 
array([[1, 2, 3, 4],
       [2, 2, 3, 4],
       [1, 4, 5, 6],
       [3, 2, 3, 4]])

In [52]: groupby_mean(a)
Out[52]: 
array([[1., 3., 4., 5.],
       [2., 2., 3., 4.],
       [3., 2., 3., 4.]])

Answer 3

Вы можете использовать панд?

import pandas as pd

df = pd.DataFrame(
    [[1, 2, 3, 4],
     [2, 2, 3, 4],
     [1, 4, 5, 6],
     [3, 2, 3, 4]], 
    columns=list('abcd')
    )

df.groupby('a').mean().reset_index().values
# returns:
array([[1, 3, 4, 5],
       [2, 2, 3, 4],
       [3, 2, 3, 4]], dtype=int64)