Блог

У меня есть код, приведенный ниже

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

Points — это массив n x d, label — это массив размером 1 x n, со значениями до k, а k — это число.

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

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

Я предполагаю, что вам нужно векторизованное решение. Вот один с 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