CODE THANKS FESTIVAL 2018 F-Coins on the tree

本番は誤読をした...。

解法

 辞書順最小を作れという指示がある。辞書順最小ということは,前の数字が小さいことが正義である。ならば,先頭から「1にできるか?」を判定し,できるなら確定,できないなら「2にできるか?」ということを判定,というように貪欲に試していこうという方針は自然である。

 それでは,試しに1番目のコインを頂点vに置けるかということを判定することを考えよう。まず,コインを頂点vに置いたのだから,木の根からの距離をd_v(ただし,d_{根}=1とする)とすると,手数をd_vだけ消費する。ここから,残りのM-1枚を,K-d_v手でピッタリ置ききる必要がある。このようなことが可能であるvの条件がわかればよい。
 まず,vに置くと決めたら,vの子孫たちはもう使えないことに注意しよう。つまり,1番目のコイン以後に使える頂点の数が減るのである。このとき,もし使える頂点の数がM-1を下回るようなことがあれば,1番目のコインをvに置くことはできない。
 使える頂点の数はクリアしたとして,手数の問題を考える。残り手数をピッタリ消費する,という条件は扱いにくいが,このような場合は可能な手数の上限と下限を考えるとうまくいくことがある。まず,手数の上限は,到達可能な頂点のうち,d_uを大きい順にM-1個取ったときの和である。消費する手数をできるだけ多くするには,できるだけ根から深いところにコインを置くようにすればよいので,このことがわかる。同様に,手数の下限は,d_uを小さい順にM-1個取ったときの和である。理由は上限と同様である。そして,この上限と下限の間の手数の消費は必ず可能である。手数の下限の状態から,手数の上限の状態までコインを順番に動かすことを考えれば,このことがわかる。以上より,K-d_vがこの上限と下限の間にあれば,コインを頂点vに置くことが可能となる。
 この判定は高速にできるだろうか。高速といっても,本問ではO(N)くらいかけてよい。それならば,到達可能な頂点をすべて訪問し,d_uのリストを作る。リストのサイズがM-1より大きいことを確認した後,これらの大きいほう,そして小さいほうからM-1の和をそれぞれとり,K-d_vと比較すればよい。sortを使ってもO(N\log N)でできる。
 1番目のコインをvに置くことが可能だとわかったら,Kd_vだけ減らし,2枚目のコインについても同様にどの頂点に置けるかを番号の小さい順に調べる。このようなことを繰り返せば,辞書順最小のコインの置き方が求まる。もし,あるコインがどの頂点にも置けないということであれば,-1を出力すればよい。

以上をまとめて,次のようにすればこの問題が解ける。

  1. 最初に,各頂点vについてd_vを求める。
  2. i番目のコインを頂点vに置けるか,ということをvが小さい順に試す。到達可能な頂点を訪問しながら,d_uのリストを作って,判定を行う。置ける条件は,①d_uのリストのサイズがM-i以上であること かつ ②d_uの小さいほうからM-i個の和をsum1,大きいほうからM-i個の和をsum2とすると,sum1 \leq K-d_v \leq sum2であること の2つ。ただし,このKi番目のコインを置こうとしている時点での残り手数を表している。頂点vに置けるとわかれば,Kd_vだけ減らして,次のコインに進む。置けないなら,次の頂点を試す。どの頂点にも置けないとなれば,-1を出力して終わり。
  3. M枚目のコインまで置けたら,結果を出力して終わり。

全体の計算量はO(MN^2\log N)である。

実装

Submission #4123308 - CODE THANKS FESTIVAL 2018
dfsが3つあるが,それぞれの役割は以下にようになる。

  • dfs(n):d_nを計算する関数。
  • dfs2(n,m):判定パートで,到達可能な点を訪問する関数。used_nというのは,それがtrueのときに頂点nにはコインを置けないということを表してる。頂点nを訪問したら判定用のリストにd_nを追加する。nの子がmまたはused_nがtrueなら,その子孫はもうコインが置けない頂点なので,その先には訪問しない。
  • dfs3(n,m):ある頂点vにコインを置くと決まった後,vの子孫にはもうコインは置けないということを記録しに行く関数。頂点番号の小さい順から試すということをするため,この作業を怠ってはいけない。

判定パートが若干混乱するが,整理すると綺麗な問題だなあと思う。