54. ライブラリのバグを疑ったけど普通に実装ミスだった話

今日は「競プロ典型 90問」の第12問を解いていました。

atcoder.jp

問題概要としては、 のマス目があり、1マス赤塗りするか、現在の状態で指定されたマスから「赤マスである箇所」だけを通って指定されたマスへ移動できるか、という判定を行うものです。

愚直に考えると「移動判定」フェーズで毎回探索を行う必要がありますが、もちろんそれでは間に合わないので工夫する必要があります。
少し考えて、「互いに移動可能な赤マスの集合」を考えれば良いと気付き、赤塗りするごとに上下左右の赤マスと自分を同じ集合に入れる、つまり Union-Find を使えば良さそうだと思い、実装しました。

ところが、何故か1問だけ WA になりました。

1問だけ…ということは何らかのコーナーケースに引っ掛かっているのでは？と思い、（自分で書いている）ライブラリにある Union-Find の実装を疑いました。
困るのは、どんなパターンでバグっているのかが分からないということですが、幸運なことに競プロ典型はテストケースが公開されています。

github.com

これで、実際に間違った答えを出した時点での盤面の様子を出力してみました。

大きい盤面なので全部は写せていませんがこんな感じです。どうも、「大きな2つの塊・境界が非常に狭い」というテストケースのようです。ここから、「これってライブラリじゃなくて普通に実装がバグっているのでは…」と疑いました。実際そのとおりでした。

以下がバグっている実装です。isRed は盤面を1次元に落とした「赤マスであるか」を表す bool 配列です。Union-Find（tree）も1次元に落としています。

// 色塗り
var r = query[1] - 1;
var c = query[2] - 1;
var idx = r * w + c;
isRed[idx] = true;

if (idx - w >= 0 && isRed[idx - w])
{
    tree.Unite(idx, idx - w);
}
if (idx + w < isRed.Length && isRed[idx + w])
{
    tree.Unite(idx, idx + w);
}
if (idx - 1 >= 0 && isRed[idx - 1])
{
    tree.Unite(idx, idx - 1);
}
if (idx + 1 < isRed.Length && isRed[idx + 1])
{
    tree.Unite(idx, idx + 1);
}

冷静に見ると、インデックスの条件式が明らかにおかしかったです。

if (idx - 1 >= 0)

意図としては「今注目している左のマスを見たいが、最左端では見ない」なのですが、2次元盤面を1次元に落としているので、これでは「左上」の盤面しか弾けていないことになります。つまり、それ以外はループして前の行の右端を判定してしまっていることになります。

同様に、最右端判定も全然見当違いのものになっています。なんでこれでテストケースほとんど通ったんだ…というレベルです。
正しい実装は以下のとおりです。

// 色塗り
var r = query[1] - 1;
var c = query[2] - 1;
var idx = r * w + c;
isRed[idx] = true;

if (r > 0 && isRed[idx - w])
{
    tree.Unite(idx, idx - w);
}
if (r < h - 1 && isRed[idx + w])
{
    tree.Unite(idx, idx + w);
}
if (c > 0 && isRed[idx - 1])
{
    tree.Unite(idx, idx - 1);
}
if (c < w - 1 && isRed[idx + 1])
{
    tree.Unite(idx, idx + 1);
}

これで普通に AC できました。

時間を溶かして悲しいですが、教訓を得たので次回に活かしていきましょう…。

• ライブラリを疑う前にまず実装を疑え
• 2次元平面を1次元に落とす必要があるなら、境界値判定に気をつける