ABC186 D - Sum of difference

備忘録

問題

atcoder.jp

回答

import sys
import os
import math
import bisect
import itertools
import collections
import heapq
import queue
import array
import time
import numpy as np
from numpy import linalg as LA

# 時々使う
# import numpy as np
# from decimal import Decimal, ROUND_HALF_UP
# from scipy.sparse.csgraph import csgraph_from_dense, floyd_warshall
# from collections import defaultdict, deque

# 再帰の制限設定
sys.setrecursionlimit(10000000)


def ii(): return int(sys.stdin.buffer.readline().rstrip())
def il(): return list(map(int, sys.stdin.buffer.readline().split()))
def fl(): return list(map(float, sys.stdin.buffer.readline().split()))
def iln(n): return [int(sys.stdin.buffer.readline().rstrip())
                    for _ in range(n)]


def iss(): return sys.stdin.buffer.readline().decode().rstrip()
def sl(): return list(map(str, sys.stdin.buffer.readline().decode().split()))
def isn(n): return [sys.stdin.buffer.readline().decode().rstrip()
                    for _ in range(n)]


def lcm(x, y): return (x * y) // math.gcd(x, y)


# MOD = 10 ** 9 + 7
MOD = 998244353
INF = float('inf')


def main():
    if os.getenv("LOCAL"):
        sys.stdin = open("input.txt", "r")

    N = ii()
    A = il()
    A.sort()
    acc = list(itertools.accumulate(A))

    ret = 0
    for i in range(1, N):
        ret += A[i] * i - acc[i-1]
    print(ret)


if __name__ == '__main__':
    main()

考え方

愚直に全てのi, jを求めるような二重のforを記述した場合、
N <= 2 * 10 ** 5なので間に合いません。
そのため、累積和を用います。

まず、|Ai - Aj|について、
これはAをソートすることでAj - Aiとすることが出来ます。
求められているのは絶対値の和なので、
計算する順番が変わっても回答は変わりません。
そのため、ソートして必ずjの方が大きい数字になるように並び替えることで、絶対値は不要になります。

次に、各和を高速に求める方法を考えます。
入力例1のようにA = [1, 2, 5]のとき、
j = 2で行われる計算は5 - 1, 5 - 2の2回です。（i = 0, 1）
このとき、Ajが出現する回数はj回、Aiとなる値は1, 2なので
予めAiを累積和で求めておくことによって、1回の計算で求めることが出来ます。
Aj * j - (Aiの累積和) = 5 * 2 - (1 + 2)
（この辺りは公式解説の動画が非常に分かりやすかったです）