ARMERIA

Rubyと競技プログラミングの話 AtCoderやCodeforcesの問題解説記事が多め。

AtCoder AGC023 B - Find Symmetries O(N^3)解法でのTLE対処

AtCoder Grand Contest 023 お疲れ様でした。A, Bの2完でした。(C通したかった…)

今回のB問題について、想定解である O(N^{3}) 解法を素直に実装すると、RubyではTLEしてしまったので、その対処法をまとめたいと思います。

問題

B - Find Symmetries

入力としてN * N の正方行列が与えられます。

N * N の正方行列を縦横にずらす方法は、縦 N 通り、横 N 通りで N * N 通りあります。(はみ出たぶんは \mod N でループします)

その N * N 通りのうち、対称行列になっているものはいくつありますか、という問題。

想定解

公式解説PDF

愚直解は、N * N 通りのずらし方それぞれについて、N * N 個の要素をチェックします。計算量は O(N^{4}) です。

これはさすがに間に合わないので、何か規則性はないか考えてみたところ…対称行列を「右下斜め」にずらしたものは、やはり対称行列である、ということを思い付きました。逆もまた然りです。

具体的には、以下の図で説明できます。行列を「右下斜め」にずらしても、比較対象のペアとなる要素は変わりません。

f:id:betrue12:20180429130213p:plain

「右下斜め」のずらし方は N 通りで、これを同一視できるので、計算量が N 1つぶん減って O(N^{3}) となります。実装においては、列をずらさずに、行だけを N 通りずらしたものをチェックすればよいでしょう。もちろん逆でも大丈夫です。

…というところまでが公式解説の内容です。C++等の高速な言語であれば、これを素直に書けば通るのではないかと思います。

Ruby実装(TLE)

コンテスト中最初に提出したコードがこちらです。TLEしています。

TLE1:Submission #2427824 - AtCoder Grand Contest 023

ほぼ同じ構造で、なるべく無駄を省こうとしたコードがこちらです。チェックする要素を約半分に削減したり、return による大域脱出を試みたりしています。これでもTLE。

TLE2:Submission #2434543 - AtCoder Grand Contest 023

これは O(N^{3}) の手続きを素直に3重ループで書いたものであり、1つ1つの要素(文字1つ)について==での比較を行っています。

制約は N\leq 300 なので、最も時間が掛かる入力ケースは a が300×300個ある行列です。これをTLE2のコードで実行すると、AtCoderのコードテストで

f:id:betrue12:20180429134345p:plain

こうなります。惜しい…。

Ruby実装(AC)

コンテスト中にACしたコードがこちら。1699ms、ギリギリ間に合っています。

AC1:Submission #2429046 - AtCoder Grand Contest 023

TLE2のように関数化して return を使うとこうなります。実行時間にほとんど差はないです。

AC2:Submission #2434803 - AtCoder Grand Contest 023

どうしたのかというと、「転置行列」を使いました。

Rubyに限った話ではないと思いますが、2重配列というのは「配列の配列」です。行列の ij 列の要素を ss[i][j] と記載できるように2重配列を構成した場合、行それぞれが配列となっています。

ss = [            # 以下の行列を示す。
    [1, 2, 3],    # 1 2 3
    [4, 5, 6],    # 4 5 6
    [7, 8, 9]     # 7 8 9
]

そこで、「要素同士の比較ではなく、行同士の比較にすれば、Rubyで処理するループは2重ループになるのではないか?」と考えました。配列同士の比較はおそらく O(N) なので全体オーダーは変わっていませんが、Rubyの内部実装(C言語)で処理されるので高速化が期待できます。定数倍の改善と言い換えても良いかもしれません。

実際に比較しなければいけないのは行と列ですが、これを行と行の比較にしたいです。ということで思いつくのは転置です。入力行列の転置行列を予め格納しておいて、それを対称行列の判定に使えないか考えてみます。

f:id:betrue12:20180429142107p:plain

横に1つずらした行列が対称行列であるための条件は、例えば1行目 [3,1,2] と1列目 [3,6,9] が等しいことです。この [3,6,9] は、転置行列の3行目と一致します。他の列についても同様に、転置行列のほうに対応する行が存在します。当たり前といえば当たり前なのですが。

この性質を使えば、「入力行列の行をシフトしたもの」と「入力行列の転置行列の行」を比較することで、対称行列のチェックができます。

# ssが元の行列、sstが転置行列
(0...N).each do |a|
    success = true
    (0...N).each do |i|
        unless ss[i] == sst[(i+a)%N]
            success = false
            break
        end
    end
    if success
        sum += N
    end
    (0...N).each do |i|
        ss[i].rotate!(1)
    end
end

行のシフトには Array#rotate! を使いました。破壊的メソッドを使って、ずらし幅が1つ進むごとに1つずつシフトしています。多くの場合、非破壊的メソッドよりも破壊的メソッドのほうが速いです。

Array#rotate! は、引数が正の数の場合配列を左にずらすので注意です(これで1WA…)

[1, 2, 3].rotate(1)    # => [2, 3, 1]

追記

Submission #2436958 - AtCoder Grand Contest 023

転置を使うというアイデアはそのままに、各行を配列ではなく文字列として扱うようにすると、==での比較がさらに高速になり120msで通りました。

さいごに

Rubyのような言語を使っていると、オーダーで示される計算量だけではなく、定数倍の部分にも気を使わなければいけないことが多いですね。時間がかかる処理をC言語実装に任せる、というのは重要な視点だなあと思いました。