京都大学プログラミングコンテスト2012 practice: D - A mul B Problem

問題記録 KUPC2012practice 行列乗算の検算 乱択アルゴリズム

https://atcoder.jp/contests/kupc2012pr/tasks/kupc2012pr_4

問題概要

N行N列の行列A, B, Cが与えられる。AB = Cであるか判定せよ。

N <= 1000

解法

普通に掛け算するとO(N3)だが、検算だけなら各要素が0か1をランダムでとるN要素のベクトルrを使って ABr = Cr かを見るというO(N2)の乱択アルゴリズムがある（詳しくは以下の記事参照）。

https://research.preferred.jp/2011/01/matrix-multiplication-and-polynomial-identity/

これをやると AB=C なら必ず ABr = Cr で、AB!=C なら1/2で ABr != Cr となるので、100回ほど回して一度でも ABr != Cr となるかどうかを見れば十分な速度・精度で答えを得られる（上の記事の「AB≠C の時は確率(1/2)kでNoと出力する」は「確率1-(1/2)kでNoと出力する」の打ち間違い？）。

感想

なるほどね

コード

import std.stdio, std.array, std.string, std.conv, std.algorithm;
import std.typecons, std.range, std.random, std.math, std.container;
import std.numeric, std.bigint, core.bitop, core.stdc.string;

void main() {
    auto N = readln.chomp.to!int;
    auto A = iota(N).map!(_ => readln.split.map!(to!long).array).array;
    auto B = iota(N).map!(_ => readln.split.map!(to!long).array).array;
    auto C = iota(N).map!(_ => readln.split.map!(to!long).array).array;

    int yes = 0, no = 0;

    foreach (_; 0..100) {
        auto v = iota(N).map!(_ => uniform(0, 2).to!long).array;
        auto t = v.dup;
        auto u = new long[](N);

        foreach (i; 0..N) foreach (j; 0..N) {
            u[i] += B[i][j] * v[j];
        }
        v = u.dup;
        u[] = 0;

        foreach (i; 0..N) foreach (j; 0..N) {
            u[i] += A[i][j] * v[j];
        }
        v = u.dup;
        u[] = 0;

        foreach (i; 0..N) foreach (j; 0..N) {
            u[i] += C[i][j] * t[j];
        }
        t = u.dup;

        (v == t ? yes : no) += 1;
    }

    writeln(no ? "NO" : "YES");
}