[TLE地獄からACへ] RustでABC405 C「Sum of Product」を O(n²)→O(n) に最適化した話

目次

1. 問題概要
2. 初手：組合せ全列挙で TLE
3. 改善1：イテレータでメモリ削減（それでも TLE）
4. 改善2：累積和で突破
5. まとめと学び

Atcoderの問題

長さNの整数列A = (A1,A2...)に対して
A1×A2＋A2×A3＋A3×A...としてすべて二つの積の和を求めるものです。
単純な二重ループは で 10^10 回越え、制限時間 2 秒に到底収まらない。

これを見て私は思いました。ふーん、簡単じゃね？って。
なお、この後。。。

初手：組合せ全列挙で TLE

方針

長さ3の整数列A = [10,20,30]としたときに、sum = 10×20＋20×30＋30×10　が答えとなります。

この問題を見たときに、コンビネーションだ！！勝ったと思いました。
コンビネーションはm個の異なるものから n個選ぶ方法（組合せ）の数のことで↓のように表せます。

{}_m \mathrm{C}_n

今回だと 3個の配列の中から2個を選んで掛け算するので

{}_3 \mathrm{C}_2

A = [10,20,30]としたときに、出力されるのは下のような形になります。

[[0,1],[1,2],[2,1]]

これらをAのインデクスとして掛け算していけば答えに行くという考えです。

すべてのコードは以下です。

use itertools::Itertools;
use proconio::input;
fn main() {
    input! {
        _n : usize,
        _a : [usize; _n]
    }
    let mut combinations = Vec::new();
    for comb in (0.._n).combinations(2){
        combinations.push(comb);
    }
    let mut sum_final = 0;
    for comb in combinations{
        let mut sum = 1;
        for c in comb{
            sum *= _a[c];
        }
        sum_final += sum;
    }
    println!("{}",sum_final);
}

結果

TLEでメモリが3.6GBほどありました。
実行時間は脅威の2510 ms

考察

問題点だらけです。組み合わせの数は2×10^10
それをvec丸ごと保存。結果無限に増えるためにメモリとコピー orz
それは遅くなるわけだ。

改善1：イテレータでメモリ削減（それでも TLE）

よく見てみると下の部分に問題があるとわかりました。

let mut combinations = Vec::new();
for comb in (0.._n).combinations(2){
    combinations.push(comb); // Vecに全組合せを保存している
}

以下二つの問題点があります。

• (0..n).combinations(2)を使って組合せの全リストを一度ベクタに格納している。
• このため、組合せの数だけメモリ確保・コピーが発生し、メモリの確保コストが非常に大きい。

結果として実質的に、メモリに対して大きな負荷がかかり、定数因子が非常に大きくなってしまっているということに気づきました。

そこで組合せを その場で逐次処理 し、Vec への退避をやめることに

use itertools::Itertools;
use proconio::input;
fn main() {
    input! {
        _n : usize,
        _a : [usize; _n]
    }
    let mut sum = 0;
    for _comb in (0.._n).combinations(2) {
        sum += _a[_comb[0]] * _a[_comb[1]];
    }

    println!("{}", sum);
}

こちらの狙いは組合せをベクタに格納せず、直接イテレータとして各組合せをその場で処理。要なメモリ確保がなくなり、定数因子が小さくなるです。

結果

メモリの結果は 8450KBで実行時間は2210 ms

考察

メモリは劇的に減りましたが、計算量は相変わらず。
2200 ms を切れず、TLE 継続。

改善2：累積和で突破

もともとの式は、

\sum_{1 \le i < j \le N} A_i \times A_j

ですが、これをよく見ると、以下のように書き換えることができます。

\sum_{i=1}^{N} A_i \left(\sum_{j=i+1}^{N} A_j\right)

つまり、各 A_i を固定したとき、その後ろ側にある全ての要素の和を求めて掛け合わせればよい、ということです。

この後ろ側の合計値を高速に計算するために用いるのが 累積和（prefix sum） というテクニックです。

累積和とは？

累積和（prefix sum）とは、元の配列 A に対して、その要素を前から順に足し合わせて新たな配列を作る手法です。

具体的には、
配列 A が以下のとき、

A = [A_1, A_2, A_3, ..., A_N]

累積配列 Sは、

S_0 = 0
S_1 = A_1
S_2 = A_1 + A_2
S_3 = A_1 + A_2 + A_3
...
S_N = A_1 + A_2 + ... + A_N

累積和を用いると、ある範囲の和が簡単に求められます。
例えば、元の配列で [l, r] 区間の合計値は次のように簡単に計算できます。

A_l + A_{l+1} + ... + A_r = S_r - S_{l-1}

\sum_{j=i+1}^{N} A_j = S_N - S_i

この考えを用いたRust実装
累積和のアイデアをそのままRustコードに落とすと、次のような非常にシンプルかつ高速なコードになります

use proconio::input;

fn main(){
    input!{
        n : usize,
        a : [usize; n]
    }
    let mut sum = 0;
    let mut prefix_sum = 0;

    for x in a {
        sum += x * prefix_sum;
        prefix_sum += x;
    }
    println!("{}", sum);
}

結果

メモリ5832 KBで実行時間は合計7msで制限時間を大幅に下回る速度で AtCoder の問題を通すことができました。

まとめと学び

1. 失敗を可視化すると改善ポイントがクリアになる。まずは動くものを出して計測すること
2. メモリをケチるだけでは限界。ボトルネックはアルゴリズムのオーダーにあったこと

お役に立てれば幸いです！！