Как заменить аргумент оси в NumPy

Question

У меня есть код как показано ниже

np.array([points[label==k].mean(axis = 0) for k in range(self.k)])

Точки - это массив n x d, метка - массив 1 x n со значениями до k, а k - число.

Моя цель - удалить аргумент оси и получить тот же результат, а также индексирование с меткой части массива == k, я хотел бы переписать.

У кого-нибудь из вас есть способ сделать это?

Answer 1

Я предполагаю, что вы ищете векторизованное решение.Вот еще один с matrix-multiplication -

def matmul(points, label):
    k = label.max()+1
    mask = label == np.arange(k)[:,None]
    out = mask.dot(points)/mask.sum(1,keepdims=True)
    return out

Вот еще один с np.add.reduceat -

def add_reduceat(points, label):
    k = label.max()+1
    sidx = label.argsort()
    ps = points[sidx]
    ls = label[sidx]

    cutidx = np.flatnonzero(np.r_[True,ls[:-1] != ls[1:],True])
    lens = np.diff(cutidx)
    out = np.full((k,points.shape[1]),np.nan)

    idx_rows = ls[cutidx[:-1]]
    mean_vals = np.add.reduceat(ps,cutidx[:-1],axis=0)/lens[:,None]
    out[idx_rows] = mean_vals
    return out

Пример выполнения -

In [220]: n,d,k = 10000,100,100
     ...: np.random.seed(0)
     ...: points = np.random.rand(n,d)
     ...: label = np.random.randint(0,k,(n))

In [221]: out0 = np.array([points[label==k_i].mean(axis = 0) for k_i in range(k)])

In [222]: np.allclose(matmul(points, label),out0)
Out[222]: True

In [223]: np.allclose(add_reduceat(points, label),out0)
Out[223]: True