vec<u32>の総和はsum<u64>()でもu32の最大値でオーバーフローする

要素がu32型のvecの総和を求める際、いつもはfor式を使い、sum()を使わないのだが、ちょっとかっこつけて使ってみたら失敗したという話です。

和を取るとき、オーバーフローするというのはよくあるのだが、sum()にをつけてもダメだった、結果をキャストするだけらしい。今日はこれでハマった。

そもそも、N=10⁹、A[i]=10⁹ (0≦i<N)みたいな、よくある設定がありますが、あれは符号なし32bit整数型(unsigned int/u32)で計算すると桁あふれすることを狙っている場合があります。

今回の計算の場合以下のように見積もると、32bitでは足りない値が出てきました。そこで最後の総和だけは64bitにしようと考えました。

• Nはlog_2(2*10⁵)=17.60…bit*1つまり、Nは18bitで表現できる。
• A[i]も同様に30bitくらい。
• ΣA[i]の最大値はN×A[i]の次元になるので、logだから上記2つを足して30bit+18bitくらい必要。

というか、わりとすぐに犯人は総和だとわかっていたのだけど、うまい対応が思いつかなかった。もちろん全部u64にするのがベストである。しかし、u32で計算できるのに、u64でやるのは悔しいじゃないか。

性能的には、並列化もしていないし、このこだわりに意味はない。メモリー使用量はu32で4840KB、全部u64でも5580KB。実行時間に至っては本当に差がない。一応vecの容量を計算してみたけど、

• 64 bit210⁵ = 1526 KiB
• 32 bit210⁵ = 763 KiB

という感じ。ただ、Rustでできるのか知らないけど、ベクトル演算ができるAVX512命令とかで上手く並列化できた場合は一部の計算(乗加算)は速度が倍になるかもしれない。

この例はAtCoder Beginner Contest 246(ABC246)のCである。 この問題の入力は≦10⁹なので、u32で格納はできる。またdkという変数はa/xなのでaを超えない。aとkはひたすら減るだけ。加算があると問題になるのだが、加算は最後まで行われない。

まずWrong Answerなコードを:

use proconio::input;

fn main() {
    input! {
        n: usize,
        mut k: u32,
        x: u32,
        mut a: [u32;n],
    };
    a.sort_unstable();
    a.reverse();
    for i in 0..n {
        let dk = a[i] / x;
        if k <= dk {
            a[i] -= k * x;
            k -= 0;
            break;
        } else {
            a[i] -= dk * x;
            k -= dk;
        }
    }
    a.sort_unstable();
    a.reverse();
    if 0 < k {
        for i in 0..n {
            if a[i] != 0 {
                a[i] = 0;
                k -= 1;
            }
            if k == 0 {
                break;
            }
        }
    }
    println!("{}", a.iter().sum::<u32>()); // オーバーフローする
}

次に、Acceptされた方を。結局キャストするので、なんかありがたみが薄いけど。しかし、このaccとして渡されていく値に蓄積していくので、これはどうしてもu64じゃないとまずい。この部分は結局u32でも、u64でも64bitで計算せざるを得ないので、並列化してもbit数削減による並列度向上はない。速くならない。

//(前略)
    println!("{}", a.iter().fold(0u64, |acc, v| acc + (*v as u64))); // オーバーフローする
}

*1:関数電卓的な計算はGoogleで: log_2(2*10^5) - Google 検索。複雑すぎてGoogleが電卓モードにならないときはlog_2(Power[10,5]) - Wolfram|Alphaで。