CODE THANKS FESTIVAL 2018 - C. Pair Distance

beta.atcoder.jp

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