【AtCoder】ABC 405 C - Sum of Product | 緑コーダーが解くAtCoder

atcoder.jp

配点: 300 点 / 実行時間制限: 2 sec / メモリ制限: 1024 MB / Difficulty: ??? / NoviSteps: 4Q

問題概要
制約
考察
- 累積和を用いる
- 式変形から導く
実装例
- 考察1
- 考察2
コメント

問題概要

長さ $N$ の整数列 $A$ が与えられるので、 $\displaystyle \sum_{1\leq i \lt j\leq N} A_iA_j$ の値を求めよ。

制約

$2\leq N \leq 3\times 10^5$
$1\leq A_i \leq 10^4$
入力は全て整数。

考察

二重シグマを愚直に計算していたのでは、 $O(N^2)$ の計算量となってしまい TLE する。

これを避けるためには、どうにかして計算量を改善しなければならない。様々な考え方ができるが、ここでは汎用的な2つの例を書くことにする。

累積和を用いる

$N = 4$ として考えてみると、求める答えは以下のように工夫して計算できる。

$\begin{align*} \displaystyle \sum_{1\leq i \lt j\leq 4} A_iA_j & = A_1 \times A_2 + A_1 \times A_3 + A_1 \times A_4 + A_2 \times A_3 + A_2 \times A_4 + A_3 \times A_4 \\ & = A_1 \times A_2 + (A_1 + A_2) \times A_3 + (A_1 + A_2 + A_3) \times A_4 \end{align*}$

括弧でくくった部分は累積和の形となり、この部分は定数時間で計算できるようになる。

$N$ を一般の場合に変えた場合のアルゴリズムは以下の通り。

答え用の変数 $\mathrm{ans}$ と、累積和用の変数 $s$ をそれぞれ $\mathrm{ans} = 0,\, s = A_1$ で初期化する。
について、以下の処理を行う。
- $\mathrm{ans} \leftarrow A_i \times s$
- $s \leftarrow s + A_i$

なお、今回は累積和テーブルの区間和が1個ずつ変更されないので、「前処理 $O(N)$ で累積和テーブルを構築する」ようなことはしていない。

オーバーフローに注意すること。

式変形から導く

高校数学では、以下のような変形をよく用いたと思う。

$\begin{align*} (a + b)^2 & = a^2 + b^2 + 2ab \\ \iff ab & = \frac{(a + b)^2 - (a^2 + b^2)}{2} \\ (a + b + c)^2 & = a^2 + b^2 + c^2 + 2(ab + bc + ca) \\ \iff ab + bc + ca & = \frac{(a + b + c)^2 - (a^2 + b^2 + c^2)}{2} \end{align*}$

これを一般化すると、以下のようになる(証明は割愛するが、帰納法で可能だと思う)。

$\begin{align*} \left( \displaystyle \sum_{i=1}^{n} a_i \right)^2 & = \sum_{i=1}^{n} a_{i}^2 + 2 \sum_{1\leq i \lt j\leq n} a_i a_j \\ \iff \sum_{1\leq i \lt j\leq n} a_i a_j & = \frac{\left( \displaystyle \sum_{i=1}^{n} a_i \right)^2 - \displaystyle \sum_{i=1}^{n} a_{i}^2}{2} \end{align*}$

あとは、これをオーバーフローに注意しつつ実装すればよい。

実装例

以下のコードは 0-indexed である。

考察1

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

#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)

using ll = long long;

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

int main()
{
    int N;
    cin >> N;
    vector<ll> A(N);
    rep(i, 0, N) cin >> A[i];

    ll ans = 0, sum = A[0];
    rep(i, 1, N)
    {
        ans += A[i] * sum;
        sum += A[i];
    }

    cout << ans << endl;

    return 0;
}

atcoder.jp

考察2

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

#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)

using ll = long long;

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

int main()
{
    int N;
    cin >> N;
    vector<ll> A(N);
    rep(i, 0, N) cin >> A[i];

    ll sum = 0, sq_sum = 0;
    rep(i, 0, N)
    {
        sum += A[i];
        sq_sum += A[i] * A[i];
    }

    ll ans = (sum * sum - sq_sum) / 2;

    cout << ans << endl;

    return 0;
}