Повышение производительности за счет исключения python для цикла при использовании уникальных значений numpy - PullRequest
1 голос
/ 14 января 2020

У меня есть два numpy массива, A с формой (N,3) и B с формой (N,), и я генерирую из вектора A вектор с уникальными записями, например:

A = np.array([[1.,2.,3.],
              [4.,5.,6.],
              [1.,2.,3.],
              [7.,8.,9.]])

B = np.array([10.,33.,15.,17.])

AUnique, directInd, inverseInd, counts = np.unique(A, 
                                             return_index = True, 
                                             return_inverse = True, 
                                             return_counts = True, 
                                             axis = 0)

Так что AUnique будет array([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]])

Затем я получу simil-вектор B, связанный с AUnique, и для каждой неуникальной строки в A я суммирую соответствующие значения B в этом векторе, то есть:

BNew = B[directInd] 

# here BNew is [10., 33.,17]

for Id in np.asarray(counts>1).nonzero()[0]: 
  BNew[Id] = np.sum(B[inverseInd == Id])

# here BNew is [25., 33.,17]

Проблема в том, что цикл for становится чрезвычайно медленным для больших N векторов (миллионов или десятков миллионов строк), и мне было интересно, есть ли способ чтобы избежать зацикливания и / или сделать код намного быстрее.

Заранее спасибо!

1 Ответ

1 голос
/ 14 января 2020

Я думаю, что вы можете делать то, что вы хотите с np.bincount

BNew = np.bincount(inverseInd, weights = B)
BNew

Out[]: array([25., 33., 17.])
...