Усреднение до значения индекса, соответствующего значению другого массива в Numpy

Question

У меня есть один массив, в котором значения должны быть усреднены до дня, который указан как значение в другом массиве. Первый массив имеет 365 дней в качестве первой оси, а второй массив соответствует определенным c юлианским датам в диапазоне от 0 до 365, из которых следует усреднить значение из первого массива.

array1.shape = (365, 375, 700)
array2.shape = (375, 700)

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

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

Answer 1

Вы можете использовать numpy .cumsum для вычисления кумулятивной суммы по axis=0, а затем взяв некоторый индекс и разделив на этот индекс, вы получите среднее значение до этого индекса.

import numpy as np


def averages(a, b):
    return a.cumsum(axis=0)[
        b.ravel(),
        np.repeat(np.arange(b.shape[0]), b.shape[1]),
        np.tile(np.arange(b.shape[1]), b.shape[0]),
    ].reshape(b.shape) / (b + 1)


a = np.arange(12).reshape(3, 2, 2)
b = np.array([[0, 1], [1, 2]])
print(a)
# [[[ 0  1]
#   [ 2  3]]

#  [[ 4  5]
#   [ 6  7]]

#  [[ 8  9]
#   [10 11]]]
print(b)
# [[0 1]
#  [1 2]]
print(averages(a, b))
# [[0. 3.]
#  [4. 7.]]