Codeforces Round #565 (Div. 3) Problem D. Recover it!

問題概要

長さ $n$ の数列 $a$ に対して, 以下の操作によって数列 $b$ を得る.

$b = a$ として初期化する.
に対して以下で定められる数字を計算し, の末尾に追加する.
- $a_i$ が素数 : $k$ は $a_i$ 番目の素数
- $a_i$ が非素数 : $k$ は $a_i$ の $a_i$ 自身より小さい最大の約数
$b$ の要素をランダムにシャッフルする

例えば, $a = \{2, 3, 4\}$ のとき, $k$ は順に $3, 5, 2$ であるから, $b$ としてたとえば $b = \{3, 4, 5, 3, 2, 2\}$ などが得られるわけです. ( $3$ は $2$ 番目の素数, $5$ は $3$ 番目の素数であることに注意します).

問題では $b$ があたえられるので, $a$ を復元してください.

問題へのリンク

制約

$1 \le n \le 2\times 10^5$
$2 \le b_i \le 2750131$ ( $2750131$ は $199999$ 番目の素数)
元の数列 $a$ の要素は $2$ 以上 $2\times 10^5$ 以下

解法

上手にマッチングしてくださいという問題です. 少し考えるとわかることですが, $a_i$ が素数のとき, 追加される要素 $k$ は必ず $a_i$ よりも大きくなり, 逆に $a_i$ が非素数のとき $k$ は $a_i$ よりも小さくなります.

そこで, まず $b$ を降順にソートし, $b_i$ が大きい順にみていくことにすれば,

$b_i$ が素数 → $b_i$ が元の数列 $a$ に属しているとすると, $b_i$ よりも大きな素数 $k$ が手順2で追加されていることになります. これは大きい順にみていることに矛盾するため, $b_i$ は $a$ に属さない要素です.
$b_i$ が非素数 → $b_i$ が元の数列 $a$ に属していないとすると, $b_i$ を他の数から生成しなければなりませんが, これは $b_i$ が一番大きいという事実に矛盾します.

ということで, 次のようなアルゴリズムを考えれば良いです.

$b$ を降順にソートする.
$b$ を大きい順にみていき, もし $b_i$ が素数なら $b_i$ に対応する素数 $m$ ( $m$ 番目の素数が $b_i$ となるような $m$ )を答えの数列に追加し, $m$ を $b$ から削除する. $b_i$ が非素数なら, $b_i$ を答えの数列に追加し, $b_i$ の最大の約数 $k$ を $b$ から削除する.
答えを出力する.

計算量についてですが, 素数の列挙はエラトステネスの篩で実装すればいいです. 問題は 2.の $b_i$ が非素数の場合に最大の約数を見つけてくるところですが, 普通に $2$ から割れるか試していき最初に割り切れる数 $j$ をみつけてきて $b_i/j$ とすればよいです. これは $O(\sqrt{b_i})$ 程度の時間がかかりますが, 毎回試しても大丈夫です. 実行時間 $4$ 秒もあるので.

コード

Submission #145007830 - Codeforces