アルゴ式の「プログラミングで学ぶ数学 (1) シグマ計算とパイ計算」では、for文による基本的な数え上げを練習問題として提供してくれています。制約が緩いのでいわゆる愚直解でも間に合うのですが、こういった問題は数学的な工夫によって計算を高速化できる場合があります。また、より難しい問題ではこのような工夫をしないと計算が間に合わない場合が多々あります。
そこで、本記事では、そのような高速化テクニックを紹介します。公式の解説にも、発展的な解法があることは触れられているのですが、その内容までは書いていません。よって本記事では、解法についての簡単な解説も交えます。なお、本記事で「発展解法」「高速解」とは、時間計算量のオーダーが解説の解法よりも少ない解法のことを指します。
日付が前後して恐縮ですが、競プロ Advent Calendar 2021に登録させていただきました。
Q6-1-1. for文とシグマ計算 (1)
愚直解
int ans = 0;
for (int i=1; i<=N; i++) ans += i;
高速解
int ans = N*(N+1)/2;
オーダー
$O(N) \rightarrow O(1)$
解説
注意点
計算の途中でNの二乗オーダーが出てくるので、Nの数が大きいとオーバーフローします。制約が大きい場合は注意してください。long long型を使う、多倍長整数を扱える言語を使う等の対策が必要になります。
Q6-1-2. for 文とシグマ計算 (2)
愚直解
int ans = 0;
for (int i=L; i<=R; i++) ans += pow(2*i-1,2);
高速解
L--;
int Lsum = 2*L*(L+1)*(2*L+1)-6*L*(L+1)+3*L;
int Rsum = 2*R*(R+1)*(2*R+1)-6*R*(R+1)+3*R;
int ans = Rsum - Lsum;
ans /= 3;
オーダー
$O(N) \rightarrow O(1)$
解説
まず、L 以上 R 以下 の区間の和は
「R 以下 の区間の和 － L 未満 の区間の和」
＝「R 以下 の区間の和 － (L-1) 以下 の区間の和」
になります。
では、R以下の区間の和を求めましょう。
そのためには、$\sum_{i=1}^N (2i - 1)^2$ を求める必要があります。これは以下のように変形できます。
\sum_{i=1}^N (2i - 1)^2 \\
=\sum_{i=1}^N (4i^2-4i+1) \\
=\sum_{i=1}^N 4i^2 - \sum_{i} 4i + \sum_{i} 1 \\
$\sum$に関係ない添字は外に出すことが出来るので、
\sum_{i=1}^N 4i^2 - \sum_{i} 4i + \sum_{i} 1 \\
= 4\sum_{i=1}^N i^2 - 4\sum_{i} i + \sum_{i} 1\\
$\sum_{i=1}^N 1$は、1 を N 回足すため$N$です。
$\sum_{i=1}^N i$が$\frac{1}{2}N(N+1)$であることは、Q6-1-1の解説で示しました。
問題は$\sum_{i=1}^N i^2$です。結論的に言うと、これは
\sum_{i=1}^N i^2 = \frac{1}{6}N(N+1)(2N+1)
になります。証明は長くなるので省きますが、興味ある方は【基本】和の公式（2乗の和）等を参考に。
結果的に、
\sum_{i=1}^N (2i - 1)^2 = \frac{1}{3}\{2N(N+1)(2N+1)-6N(N+1)+3N\}
になります。後はこれをプログラムに写すだけです。
注意点
計算の途中でNの三乗オーダーが出てくるので、二乗の場合よりもオーバーフローの制約が厳しいです。
また、計算が多項式になっている場合は、割り算は最後にまとめてやった方が無難です。理由は、割り算を個別にやってしまうと、整数計算の切り捨ての関係で答えがおかしくなってしまう場合があるためです。
逆に、オーバーフローを気にする必要がある場合は、「まとめて割り算」はできずに、こまめに計算をする必要があります。そのあたりは逆元とも絡んできて、本記事の範囲を超えるので、興味ある方は「1000000007 で割ったあまり」の求め方を総特集！ 〜 逆元から離散対数まで 〜等を参考に。
Q6-1-3. for 文とシグマ計算 (3)
愚直解
int ans = 0;
for (int i=0; i<=N-1; i++) {
for (int j=0; j<=M-1; j++) {
ans += A[i]+B[j];
}
}
発展解法
int ans = 0;
for (int i=0; i<=N-1; i++) {
ans += A[i]*M;
}
for (int j=0; j<=M-1; j++){
ans += B[j]*N;
}
オーダー
$O(NM) \rightarrow O(N+M)$
解説
問題の「入力例1」の計算過程を見ると、A 側の数字（青丸）は M 個ずつ、B 側の数字（黄丸）は N 個ずつ使われていることがわかります。よって、それぞれを個別に計算することができます。
これは式変形によっても導出することができます。
\sum_{i}^N\sum_{j}^M(A_i+B_j) \\
= \sum_{i}^N\sum_{j}^M A_i + \sum_{i}^N\sum_{j}^M B_j \\
= \sum_{i}^N A_i \sum_{j}^M 1 + \sum_{j}^M B_j \sum_{i}^N 1 \\
= \sum_{i}^N A_i M + \sum_{j}^M B_j N \\
= M \sum_{i}^N A_i + N \sum_{j}^M B_j
Q6-1-4. for 文とシグマ計算 (4)
愚直解
int ans = 0;
for (int i=1; i<=N-1; i++) {
for (int j=i+1; j<=N; j++) {
ans += i*j;
}
}
発展解法
int ans = 6*N*N*(N+1)*(N-1)-3*N*N*(N-1)*(N-1)-2*N*(N-1)*(2*N-1);
ans /= 24;
オーダー
$O(N^2) \rightarrow O(1)$
解説
シグマのシグマで頭がこんがらがりそうになりますが、添字を整理すると以下のような形になります。
\sum_{i=1}^{N-1} \sum_{j=i+1}^N ij \\
= \sum_{i=1}^{N-1} i\sum_{j=i+1}^N j \\
= \sum_{i=1}^{N-1}(i \sum_{j=1}^N j - i \sum_{j=1}^i j) \\
= \sum_{i=1}^{N-1}\{i \frac{1}{2}N(N+1) - i \frac{1}{2}i(i+1)\} \\
= \sum_{i=1}^{N-1}\{i \frac{1}{2}N(N+1) - \frac{1}{2}i^2(i+1)\} \\
= \frac{1}{2}N(N+1)\sum_{i=1}^{N-1}i - \frac{1}{2}\sum_{i=1}^{N-1}(i^3+i^2) \\
= \frac{1}{2}N(N+1)\frac{1}{2}N(N-1)-\frac{1}{2}\{\frac{1}{4}N^2(N-1)^2+\frac{1}{6}N(N-1)(2N-1)\} \\
= \frac{1}{4}N^2(N+1)(N-1)-\frac{1}{8}N^2(N-1)^2-\frac{1}{12}N(N-1)(2N-1) \\
= \frac{1}{24}\{6N^2(N+1)(N-1)-3N^2(N-1)^2-2N(N-1)(2N-1)\} \\
なお、計算過程で以下の公式（三乗和）を用いています。
\sum_{i=1}^N i^3 = \frac{1}{4}N^2(N+1)^2
別解（追記）
int sum = n*(n+1)/2;
int dsum = n*(n+1)*(2*n+1)/6;
int ans = (sum*sum-dsum)/2;
上記のように考えることもできます。
Q6-1-5. for 文と Π 計算
問題リンク
これについては、公式の解が（たぶん）最適解となります。