FFTを用いた解法を紹介します。

まず色を \(0,1,2\) の番号に対応させます。また、すべて色 \(0\) になるような斜め線の数を求めることにします。(他の色についても同じように試せばよいです。)

以下のような関数が定義できると嬉しいです。(嬉しい点については後述します)

• \(f(x,y)=1\) (色 \(x\) と色 \(y\) で塗った時に色 \(0\) になるとき)
• \(f(x,y)\ne 1\) (そうでないとき)

これは、\(d\) を\(2\) 以上の適当な正整数として、色 \(0\) を \(1\) に、色 \(1\) を \(d\) に、色 \(2\) を \(\frac{1}{d}\) に対応させて積を求めればよいです。これが全ての要件を満たしていることは簡単に確認できます。

よって、このように数字を対応させて適切に畳み込み、数がその斜め線にあるマスの個数と一致する斜め線が答えの候補です。

答えでない斜め線を答えの候補に入れてしまう確率が \(0\) に近ければ、このアルゴリズムは正当と言えます。

実数範囲で行うFFTでは、誤差を考えないことにすれば \(d\) を十分大きく取ることで誤る確率を \(0\) にできます。\(d\) または \(d^2\) が加算された時点で大きさをオーバーし、\(\frac{1}{d}\) または \(\frac{1}{d^2}\) が加算された時点で整数にすることが不可能になるためです。\(\bmod\) をとるNTTで確率解析する方法はよくわからないのですが、こちらでも \(d\) を大きめに取ればACすることができました。