AtCoder-ABC186 D - Sum of difference【Python解答例】 - ebisukeプログラミング初心者脱出黙示録

AtCoderのパナソニックプログラミングコンテスト（AtCoder Beginner Contest186）のD問題について、コンテストでは解けなかった問題を勉強していきます。
Panasonic Programming Contest (AtCoder Beginner Contest 186) - AtCoder

AtCoder Beginner Contest186 D - Sum of difference

D - Sum of difference

問題文

N個の整数A1,…,ANが与えられます。
1 ≤ i < j ≤ Nを満たす全てのi,jの組についての|Ai−Aj|の和を求めてください。
すなわち、N−1∑i=1 N∑j=i+1 |Ai−Aj|を求めてください。

制約

・2≤N≤2×10^5
・|Ai|≤10^8
・Aiは整数である。

解答例

n = int(input())

a = list(map(int, input().strip().split()))
a.sort()

ans = 0
sum_a = a[0]
for j in range(1,n):
    ans += j * a[j] - sum_a
    sum_a += a[j]

print(ans)

解説

youtubeの公式解説で勉強した内容です。

整数のリストaは入力を受け取ったあとでソートしています。
これによりAj - Aiは正の整数となるので絶対値をとる必要がなくなります。

ここから先の計算は、問題文通り2重ループとするとTLEになります（NG例として載せます）ので、jの1重ループを回します。
ans変数はAj - Aiを足し合わせた解答となる変数です。
sum_a 変数はa[j-1]までの累積和です。
j = 3のとき、[i, j]の組み合わせは[1, 3]と[2, 3]の2通り考えられます。
よってAj - Aiにおいて、Aj = a[j]×jとなりAiはj-1までの累積和になります。
この計算をNまで繰り返すことでans変数に求める値が格納されました。

NG例

n = int(input())

a = list(map(int, input().strip().split()))

ans = 0
for i in range(n-1):
    for j in range(i+1,n):
        ans += abs(a[i] - a[j])

print(ans)