2020-10-09

AtCoder ABC 177 C - Sum of product of pairs (灰色, 300 点)

AtCoder AtCoder300点 ABC-C 数え上げ問題式変形灰色diff 差分更新解空間:O(N^2)通りの選択肢数学(代数) 累積和

茶色 diff にはなると思ったけど、灰色 diff だった...「/2」が必要と感じて詰まる人も多いと思ったのに...

問題へのリンク

問題概要

$N$ 個の整数 $A_{1}, \dots, A_{N}$ が与えられる。

$1 \le i \lt j \le N$ を満たすすべての $(i, j)$ に対しての $A_{i} \times A_{j}$ の総和を 1000000007 で割ったあまりを求めよ。

制約

$2 \le N \le 2 \times 10^{5}$

考えたこと

まず、すべてのペアに対する総和を愚直に計算しようと思ったら、次のような for 文を書くと思う。

long long res = 0;
for (int i = 0; i < N; ++i) {
    for (int j = i + 1; j < N; ++j) {
         res += A[i] * A[j] % 1000000007;
         res %= 1000000007;
    }
}

しかしこれでは $O(N^{2})$ の計算量となって間に合わないので、工夫が必要となる。

このとき、受験数学に慣れている人なら、ノータイムで次のような式変形をすると思う (受験数学に慣れていなくても解ける方法もある！)。

$(求める値) = \frac{1}{2} ((A_{1} + \dots + A_{N})^{2} - (A_{1}^{2} + \dots + A_{N}^{2}) )$

よって、modint が十分に整備されていれば、この式を適用するだけになる。modint の整備が不十分の場合、mod 1000000007 の世界で「2 で割る」ということがどういうことなのかの理解が必要となる。そこで、次の 2 つの道を選ぼう。

「2 で割る」ということをする必要のない解法を編み出す
mod の世界で「2 で割る」のをどうしたらいいかを理解して、実装する (または modint で殴る)

解法 (1)：まとめられるところを上手にまとめる

一般の $N$ では考えづらいので、 $N = 4$ あたりで考えてみよう。このとき、求める値は次の合計になる (0-indexed で考える)。

A[0]*A[1]  A[0]×A[2]  A[0]×A[3]
           A[1]×A[2]  A[1]×A[3]
                      A[2]×A[3]

これは、各列 (縦方向に見る) に見ることで、次のように整理できる。次の 3 つの値の合計となる。

(A[0]) × A[1]
(A[0] + A[1]) × A[2]
(A[0] + A[1] + A[2]) × A[3]

よって、i = 0, 1, ... に対して次の値が計算できれば、 $O(N)$ で求められることがわかった！！

A[0] + A[1] + ... + A[ i ]

この値は「累積和」そのものなので、累積和を求めることで計算できる。また、直接累積和を求めなくても、次のようにしても OK！

long long res = 0;
long long sum = A[0];
for (int i = 1; i < N; ++i) {
    res += sum * A[i];
    sum += A[i];
}

これで計算量は $O(N)$ となる。

コード

ここでは modint は用いずに書いた。

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

const int MOD = 1000000007;
int main() {
    int N;
    cin >> N;
    vector<long long> A(N);
    for (int i = 0; i < N; ++i) cin >> A[i];

    long long res = 0, sum = A[0];
    for (int i = 1; i < N; ++i) {
        res += sum * A[i], res %= MOD;
        sum += A[i], sum %= MOD;
    }
    cout << res << endl;
}

解法 (2)：「2 で割る」処理を敢行する

受験数学ではおなじみの式変形

$(求める値) = \frac{1}{2} ((A_{1} + \dots + A_{N})^{2} - (A_{1}^{2} + \dots + A_{N}^{2}) )$

を活用しよう。modint で殴ると次のように書ける。計算量はやはり $O(N)$ となる。なお、modint の理論については次の記事で解説した。

qiita.com

コード

下のコードでは、数列 A 自体を vector<mint> 型で宣言して入力を受け取っている。

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

// modint
template<int MOD> struct Fp {
    long long val;
    constexpr Fp(long long v = 0) noexcept : val(v % MOD) {
        if (val < 0) val += MOD;
    }
    constexpr int getmod() const { return MOD; }
    constexpr Fp operator - () const noexcept {
        return val ? MOD - val : 0;
    }
    constexpr Fp operator + (const Fp& r) const noexcept { return Fp(*this) += r; }
    constexpr Fp operator - (const Fp& r) const noexcept { return Fp(*this) -= r; }
    constexpr Fp operator * (const Fp& r) const noexcept { return Fp(*this) *= r; }
    constexpr Fp operator / (const Fp& r) const noexcept { return Fp(*this) /= r; }
    constexpr Fp& operator += (const Fp& r) noexcept {
        val += r.val;
        if (val >= MOD) val -= MOD;
        return *this;
    }
    constexpr Fp& operator -= (const Fp& r) noexcept {
        val -= r.val;
        if (val < 0) val += MOD;
        return *this;
    }
    constexpr Fp& operator *= (const Fp& r) noexcept {
        val = val * r.val % MOD;
        return *this;
    }
    constexpr Fp& operator /= (const Fp& r) noexcept {
        long long a = r.val, b = MOD, u = 1, v = 0;
        while (b) {
            long long t = a / b;
            a -= t * b, swap(a, b);
            u -= t * v, swap(u, v);
        }
        val = val * u % MOD;
        if (val < 0) val += MOD;
        return *this;
    }
    constexpr bool operator == (const Fp& r) const noexcept {
        return this->val == r.val;
    }
    constexpr bool operator != (const Fp& r) const noexcept {
        return this->val != r.val;
    }
    friend constexpr istream& operator >> (istream& is, Fp<MOD>& x) noexcept {
        is >> x.val;
        x.val %= MOD;
        if (x.val < 0) x.val += MOD;
        return is;
    }
    friend constexpr ostream& operator << (ostream& os, const Fp<MOD>& x) noexcept {
        return os << x.val;
    }
};

const int MOD = 1000000007;
using mint = Fp<MOD>;

int main() {
    int N;
    cin >> N;
    vector<mint> A(N);
    for (int i = 0; i < N; ++i) cin >> A[i];

    mint one = 0, two = 0;
    for (int i = 0; i < N; ++i) {
        one += A[i];
        two += A[i] * A[i];
    }
    cout << (one * one - two) / 2 << endl;
}