Ответ 1
Я могу предложить решение O (n log n) для начала: пусть f i будет i-м номером результата. Мы имеем:
Проходя по массиву слева направо и поддерживая двоичное дерево поиска элементов 0 до i-1, мы можем решить все части формула в O (log n):
- Сохраняйте размеры поддерева, чтобы подсчитать элементы, большие или меньшие, чем заданные.
- Сохраняйте кумулятивные суммы поддеревьев, чтобы ответить на сумма запросов для элементов, больших или меньших, чем заданный
Мы можем заменить расширенное дерево поиска на некоторые более простые структуры данных, если мы хотим избежать стоимости реализации:
- Сортировка массива заранее. Назначьте каждое число своим рангом в отсортированном порядке
- Хранить двоичное индексированное дерево значений 0/1 для вычисления количества элементов, меньших заданного значения
- Сохранять другое двоичное индексированное дерево значений массива для вычисления сумм элементов, меньших заданного значения
TBH Я не думаю, что это можно решить в O (n) в общем случае. По крайней мере, вам нужно будет сортировать числа в какой-то момент. Но, возможно, числа ограничены или у вас есть другие ограничения, поэтому вы можете реализовать операции суммирования и подсчета в O (1).
Реализация:
# binary-indexed tree, allows point updates and prefix sum queries
class Fenwick:
def __init__(self, n):
self.tree = [0]*(n+1)
self.n = n
def update_point(self, i, val): # O(log n)
i += 1
while i <= self.n:
self.tree[i] += val
i += i & -i
def read_prefix(self, i): # O(log n)
i += 1
sum = 0
while i > 0:
sum += self.tree[i]
i -= i & -i
return sum
def solve(a):
rank = { v : i for i, v in enumerate(sorted(a)) }
res = []
counts, sums = Fenwick(len(a)), Fenwick(len(a))
total_sum = 0
for i, x in enumerate(a):
r = rank[x]
num_smaller = counts.read_prefix(r)
sum_smaller = sums.read_prefix(r)
res.append(total_sum - 2*sum_smaller + x * (2*num_smaller - i))
counts.update_point(r, 1)
sums.update_point(r, x)
total_sum += x
return res
print(solve([3,5,6,7,1])) # [0, 2, 4, 7, 17]
print(solve([2,0,1])) # [0, 2, 2]