【AtCoder】ABC 353 C - Sigma Problem | 茶コーダーが解くAtCoder

実行時間制限: 2 sec / メモリ制限: 1024 MB / Difficulty: 795

問題概要

正整数 $x, y$ に対して $f(x, y) := (x+y) \mod 10^8$ とする。長さ $N$ の正整数列 $A$ が与えられるとき、 $\displaystyle \sum_{i=1}^{N-1} \sum_{j=i+1}^{N} f(A_i, A_j)$ の値を求めよ。

制約

入力はすべて整数。
$2 \leq N \leq 3 \times 10^5$
$1 \leq A_i \leq 10^8$

考察

$N$ の最大値が $3 \times 10^5$ と大きいので、愚直に2重ループを回していては間に合わない。何か工夫が必要である。

そこで、 $1 \leq A_i \leq 10^8$ であることから $A_i + A_j$ を $10^8$ で割った商は常に $0, 1$ のいずれかになることに注意して、 $f(x, y)$ を次のように言い換えてみる。

$f(x, y) = \left\{ \begin{array}{} x + y & (x + y < 10^8) \\ x + y - 10^8 & (x + y \geq 10^8) \end{array} \right.$

こうすることで、求める式の値は $A_i + A_j \geq 10^8 (1 \leq i < j \leq N)$ となる組 $(i, j)$ の個数を $k$ として、 $\displaystyle \sum_{i=1}^{N-1} \sum_{j=i+1}^{N} (A_i + A_j) - k \times 10^8$ に等しいことがわかる。

第1項については、各 $A_i$ が $N - 1$ 回ずつ加算されることを考えると、 $\displaystyle (N-1) \sum_{i=1}^{N-1} A_i$ として求められる。

問題は第2項だ。なんとかして $k$ を $O(N)$ 未満で求めたい。そこで、こんなことを考えよう。

「昇順ソートされた $A$ について、各 $i$ に対して $A_j \geq 10^8 - A_i$ となる $j$ を数えていけばいいのでは？」

$j_0$ は二分探索 ( C++ ではlower_bound) で求めることができるので、計算量は $O(N \log N)$ 程度まで抑えられる。

あとはこれを実装するだけ...と言いたいところだが、 $k$ は $(i, j)$ の組みあわせの個数なので、二重にカウントしないように二分探索の探索範囲を $i$ に応じて変える必要があることに注意しよう。また、int型だとオーバーフローする (1敗) 。

コード

#include <bits/stdc++.h>
using namespace std;

using ll = long long;

#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)
#define all(x) (x).begin(), (x).end()

// ======================================== //

ll MOD = 100000000;

int main()
{
    ll N;
    cin >> N;
    vector<ll> A(N);
    ll sum = 0;
    rep(i, 0, N)
    {
        cin >> A[i];
        sum += A[i];
    }
    sum *= (N - 1);

    sort(all(A));

    ll k = 0;
    rep(i, 0, N)
    {
        auto j0 = lower_bound(A.begin() + i + 1, A.end(), MOD - A[i]);
        k += A.end() - j0;
    }

    cout << sum - MOD * k << endl;
}