AtCoder Beginner Contest 177を見直すC - Sum of product of pairs

C - Sum of product of pairs

問題

N個の整数が与えられます。その中から2つを取り出し掛け算を行います。全ての組み合わせの和をmod(10^9+7)で求めなさい。

mod(10^9+7)とは

コンテスト中もお世話になったんですが、以下のページがわかりやすかったです。（理解したとは言っていない）今回理解しておくのは、足し算、掛け算に関しては、計算の前後でMODしても、計算の後にまとめてMODしても答えは変わらないという点だと思います。

「1000000007 で割ったあまり」の求め方を総特集！ 〜 逆元から離散対数まで 〜 - Qiita

最初に提出したコード

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using vi = vector<int>;

int main() {
   int n;
   cin >> n;
   vi a(n);
   rep(i, n) cin >> a.at(i);
   ll ans = 0;
   for(int i = 0; i < n; i++){
       for(int j = i; j < n; j++){
           if(i == j) continue;
           ans += (ll)a.at(i)*a.at(j);
           ans %= 1000000007;
       }
   }
   cout << ans << endl;
   return 0;
}

全検索してみたんですが、TLEが出てしまいました。計算量は多分O((n-1)!)なので、入力される数列の制約が最大10^5って段階で計算量から考えてTLEっぽいなと思える様になりたい。どうにか計算量を減らさないといけないので、式変形してみることにしました、入力が[A, B, C, D]の場合、上記を式で表すと以下の様になります。

A*B + A*C + A*D + B*C + B*D + C*D

さらに項でまとめると以下の様になります

A(B + C + D) + B(C + D) + C*D

さらに、さらに書き換えると以下の様になります。

  A((A + B + C + D) - A)
+ B((A + B + C + D) - (A + B))
+ C((A + B + C + D) - (A + B + C))

長くなってますが、ここまでくるとプログラムに落とし込むのは簡単です。

上記を実装したコード

#include <bits/stdc++.h>
using namespace std;
#define rep(i, n) for (int i = 0; i < (int)(n); i++)
using ll = long long;
using vi = vector<int>;

int main() {
   int n;
   cin >> n;
   vi a(n);
   rep(i, n) cin >> a.at(i);
   ll ans = 0;
   ll sum = 0;
   rep(i, n){
       sum += a.at(i);
       sum %= 1000000007;
   } 
   for(int i = 0; i < n; i++){
       sum -= a.at(i);
       if(sum < 0) sum += 1000000007;
       ans += (ll)a.at(i)*sum;
       ans %= 1000000007;
   }
   cout << ans << endl;
   return 0;
}

引き算をしているので、0以下になった場合は、10^9+7を足す様にしていますが、その他は上に書いた式をそのままコードにした感じです。

ちなみに解説動画でカッコいい解き方が紹介されていましたが、「逆元」が使えないと解答できない様です。

「逆元」については解説が無いので、興味がある方は自分で調べてね。ってことだと思うんですが、とりあえずは「存在を知った」程度にしておこうと思います。