G - 1 Bird on Cherry Blossoms

この問題は素数判定とDFSを問います。今回は時間制限に厳しくなるような制約にしております。

各 $k$ 毎にDFSで処理するとなると、$O(N^2(N+M))$ ※$_1$で制限時間に間に合いません。なので、前処理をしておく必要があります。

よくよく考えてみると、小鳥がどこの頂点に舞い降りようと、移動できる範囲は連結成分に属する頂点間のみであることがわかります。よって連結成分ごとに集める素数の種類数を求めることで実装が可能です※$_2$。素数の種類数といいますが、$W_i$ が異なるので単に個数だけを調べるとよいです。

$k=0,1,2,...N$ に対する答えは、連結成分ごとに数えた素数の種類数(個数)の最大値 $P$ のみに依存します。なぜなら小鳥は素数の種類数がより多い連結成分内に属する頂点に舞い降りれば、必ず達成できるためです( $P\ge k$ だとしても小鳥は素数を $k$ 種類集めた時点で行動を終了することで達成できる)。$P\lt k$ なら、どの頂点に舞い降りようとしても絶対に達成できません。

しかしどの連結成分に対しても頂点に置かれた整数がすべて素数( すなわちすべての $i\ (1\le i\le N)$ に対し $W_i$ が素数 )なら、$k=0$ のときは絶対に達成することはできません。小鳥が舞い降りた頂点の整数を必ず取るためです(ある $i$ に対し $W_i$ が合成数なら、合成数が置いてある頂点に舞い降りた時点で行動を終えることで達成できる)。

以上で問題を解く方針はできました。しかしこの問題では制約を厳しくしているため、探索する頂点に置いてある整数が素数かどうか、愚直に判定していると $O(\max(\sqrt{W_i})(N+M))$ で間に合いません。それを改善するためには、事前に素数判定を行っておく必要があります。

以上から $O(N\max(\sqrt{W_i})+M)$ でこの問題を解くことができます。以下は解答例(C++,Python)です。

※$_1$ $N^2$ がつく理由は、連結成分を探索するために最悪で $O(N)$ かかるためです。さらにその判定を $O(N)$ だけ回す必要があるので、$O(N^2)$ だけかかってしまいます。

※$_2$連結成分ごと数えるときはDFSの実装によって帰りがけの順か、行きがけの順で調べることがありますが、これはどっちの順番でも正解できます。
解答例では帰りがけの順で探索しています。