Агрегировать элементы на основе вектора положения

Question

Я пытаюсь векторизовать очень простую операцию, но не могу понять, как.

Учитывая очень большой числовой вектор (более 1М позиций) и другой массив размера n с заданным набором.позиций, я хотел бы получить вектор размером n с элементами, представляющими собой средние значения первого вектора, как указано вторым

a = np.array([1,2,3,4,5,6,7])
b = np.array([[0,1],[2],[3,5],[4,6]])

c = [1.5,3,5,6]

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

Решение Vanilla Python:

import numpy as np
import time

a = np.array([1,2,3,4,5,6,7])
b = np.array([[0,1],[2],[3,5],[4,6]])

begin = time.time()

for i in range(100000):

    c = []

    for d in b:
        c.append(np.mean(a[d]))

print(time.time() - begin, c)
# 3.7529971599578857 [1.5, 3.0, 5.0, 6.0]

Answer 1

Я не уверен, что это обязательно быстрее, но вы можете попробовать:

import numpy as np

a = np.array([1, 2, 3, 4, 5, 6, 7])
b = np.array([[0, 1], [2], [3, 5], [4, 6]])

# Get the length of each subset of indices
lens = np.fromiter((len(bi) for bi in b), count=len(b), dtype=np.int32)
# Compute reduction indices
reduce_idx = np.roll(np.cumsum(lens), 1)
reduce_idx[0] = 0
# Make flattened array of index lists
idx = np.fromiter((i for bi in b for i in bi), count=lens.sum(), dtype=np.int32)
# Reorder according to indices
a2 = a[idx]
# Sum reordered array at reduction indices and divide by number of indices
c = np.add.reduceat(a2, reduce_idx) / lens
print(c)
# [1.5 3.  5.  6. ]