from itertools import combinations def slow_solve(N, Xs): ans = 0 for pairs in combinations(Xs, 2): ans += abs(pairs[0] - pairs[1]) return ans
まずは遅く総当たり。 で、高速化を考える。
まず、Xはソートしても一般性を失わない。
x0 - x1 - x2 - x3
ここで、まとめられる計算がないかを考える。 各間は何回も同じ場所を足すのでここはまとめられそう。
x1 - x2の間をd1としたとき、 ここはx0 - x2 / x0 - x3 / x1 - x2 / x1 - x3で通る場所である。
つまりこれを一般化すると、diは、x0 ~ xiまでの一つとx(i+1) ~ xnまでの一つ選ぶ時の通り数だけ通る場所であるので この回数をかけながら距離を計算すれば計算量O(N)で、カウントできることになる。 0 ~ iまでの一つの選び方はi + 1通り、 i + 1 ~ nまでの一つの選び方は n - (i + 1) 通り
def solve(N, Xs): ans = 0 Xs = sorted(Xs) for i, j in zip(range(N), range(1, N)): d = abs(Xs[j] - Xs[i]) c = ((i+1) * (N-i-1)) ans += d * c return ans