import numpy
A = numpy.array([[5.0, 5.0, 5.0], [8.0, 9.0, 9.0]])
B = numpy.array([[1, 1, 2], [3, 2, 1]])
print(numpy.bincount(B.ravel(), weights=A.ravel()))
Поиск элементов в одном массиве на основе второго массива