Senの競技プログラミング備忘録

こけた問題を自分用の解説で載せる。けんちょんさんのブログを目指したい。質的にも量的にも。こけた問題だけに限定するけど

ACL Beginner Contest F - Heights and Pairs

良問 数え上げ 畳み込み マージテク 数学 難問 要復習

https://atcoder.jp/contests/abl/tasks/abl_f

問題概要

$ 2N $人の人がいて、それぞれ身長が$ h[i] $である。 この人たちで、ペアを$ N $組作っていく。(任意の人は1つのペアにしか属してはいけない)

ここで、すべてのペアのうちの2人で、お互いの身長が異なるという前提で作ることにする。

この時、何通りのペアの組方があるのでしょう？$ 998244353 $で割ったあまりを答えてください。

ただし、別にカウントする条件は、すでにあるどの組方とでも、少なくとも1ペア以上で構成員が異なるである。

入力：$ 1 \leq N \leq 50000, 1 \leq h[i] \leq 100000 $

考え方

  1. 同じ身長の人をまとめて、DPをしてみる。$ dp[i][j] := i番目の身長グループまで見て、j人まだペアを組めてない $時の組方の場合の数
  2. 先程のDPを計算してみると、計算量は$ O(N ^ 3) $ととても厳しい。
  3. この回はライブラリゲーの問題中心だからメタ読みしてNTTどうせするんだろと思い、NTTやセグ木による高速化を試すも、この遷移ではいずれも不可能。
  4. 仕方ないので方針転換。形が包除っぽいので包除らしく攻めてみる。
  5. 包除の条件の取り方を試行錯誤してみる。何も考えずに、「身長グループiとjとkでは同じ身長のペアがあって、他は任意」のような条件で考えると、考えるべき包除の条件の数が$ O(2 ^ {O(n)}) $からどうしても落とせなくなる。
  6. うまい包除の条件 「全ての身長グループのうち、同じ身長のペアが最低$ i $ペアあるとき(それ以外は任意)」というのを見つけてそこから考える。
  7. これは本質的には 「すべての身長グループで、同じ身長のペアを$ i $ペア選ぶときの選び方」と同じ。これの漸化式を立ててみる。
  8. これの漸化式はうまく畳み込みができる！
  9. 畳み込み計算の順序に注意して(マージテク)、畳み込み演算をすれば全体で$ O(n (\log n) ^ 2) $となるので、AC。

考察

ダメな例

snukeさんが解説放送の中で提示したダメな例です。思いっきり引っかかった

身長を高さ6の人は10人、などの形で身長グループごとに分類する。以後身長グループの個数を$ n $とする。(計算するのは$ 2N $人のマッチングであることに注意)

箱根DPのように、DPの添え字に 「今まで見てきてペアをまだ作れてない人の数」を持たせて、場合の数を計算する。詳細は以下の通り。 $$ \begin{aligned} &num[i] := 身長グループiの人数 \\ &dp[i][j] := 身長グループiまで見て、その時j人まだペアを作ってないときの場合の数 \\ &dp[0][0] = 0 \\ &dp[i + 1][j + num[i] - 2 * k] += dp[i][j] \times _ {j} C _ {k} \times _ {num[i]} C _ {k} \times k!\\ &Ans = dp[n][0] \end{aligned} $$

遷移の式では、$ j $人すでに余っているときに、身長グループ$ i $では$ num[i] $人いる。この新しく加わった人と、既存の$ j $人の間で、$ k $ペアを新しく作るときを示している。 既存の$ j $人は必ず身長がグループ$ i $の人と異なるので問題はない。遷移の係数は以下の意味を持つ。なお、$ k \leq \min(j, num[i]) $を当然満たす。

  • $ _ {j} C _ {k} $ は$ j $人のうち、ペアを組む$ k $人を選ぶ。
  • $ _ {num[i]} C _ {k} $ は$ num[i] $人のうち、ペアを組む$ K $を選ぶ。
  • $ k! $は、お互いの$ k $人ずつでペアを組む際に、グループ$ i $の方を固定して相方を考えると、この形となる。

ここまで考察してよさげなDPが生まれたので、あっさりACしていただきたいところであるが、このDPの時間計算量は$ O(N ^ 3) $である。(グループ数の$ n $は$ N $に依存し、遷移にも$ O(N) $かかるため)これをどうにかして、計算量を畳み込みのNTTで落とせるのだろうか？(セグ木遷移自体が総当たり的なので無理)

しかし、これはできません。とても悲しい事実ですが、このDPの遷移をNTTで高速化することはできません。

畳み込みの遷移式は以下のとおりである。 $$ c[i + j] += a[i] \times b[j] $$

この時、$ i $が変数である$ a[i] $側と$ j $が変数である$ b[j] $側にはっきりと分けられるというのがポイントである。逆に言えば、分けられない場合は畳み込みではないということである(ただの遷移が総当たりのDP)

今回の場合、$ j $と$ k $に綺麗に分割できればいいのだが、$ _ {j} C _ {k} $のように、明らかに$ j $かつ$ k $の関数が遷移の係数にかかる。 この場合では、畳み込みを用いて高速化ができないので、この考察筋は諦めるしかない。

例え正しく立式できても、それがオレンジ中位のこの問題の即解決には綱らがないのである。世間は非情。

包除原理

では失意の気持ちを抱きながら、問題をはじめから考察してみる。すると、 「任意の～」という包除原理御用達のキーワードが眠ってるので、包除原理で攻められないかを考えてみる。

包除原理系数え上げでは、条件数を$ a $とすると、何もしなければ合計で$ O(2 ^ a) $個の条件になってしまい、簡単な問題以外では明らかにTLEする。 多くの問題では、ここからさらに踏み込んで以下のことを要求してくる。

  • $ 2 ^ a $この条件のうち、まとめて計算できるものをまとめて計算する
  • そもそも条件数を1つとかにする($ O(a) $個の計算になる)

この問題の包除原理の適用例をいくつか考えてみよう。

いくつかの身長グループで$ i $ペアの同じ身長のペアがあり、他は任意

まず考え付くのはこのタイプである。しかし、これは明らかに$ 2 ^ a $個の条件があるタイプである。身長グループを同じサイズならまとめて計算したいと考えることもできるが、$ \sum ^ {x} _ {i = 1} i = \frac{x(x + 1)}{2} $から、一応身長グループを$ O(\sqrt(N)) $種類まで抑えることができる。しかし、結局全通り試さないといけないのは変わっておらず、指数時間となってしまうためこの分け方は不適である。

すべての身長グループで$ i $ペアの同じ身長のペアがあり、他は任意

さきほどの問題は、各身長グループごとに分けていたから場合分けが指数的であったことである。ここで、 「そもそも条件数を１つとかにする」を試してみる。

すると、これはうまく行きそうだとわかる。この場合は条件数は1つなので、$ O(N) $この条件で済む。

というわけで、「すべての身長グループで$ i $ペアの同じ身長のペアがあり、他は任意」を包除の条件として攻めてみる。

包除の条件ごとの詳しい計算

ここで、$ a $人いて、任意にペアを組む場合の数は、二重階乗 を用いると、$ (a - 1)!! $となる。まだ組めてない人を1人選び、その相手を残りの人の中から選ぶ、を繰り返せばよい。

すると、「すべての身長グループで$ i $ペアの同じ身長のペアがあり、他は任意」条件の計算に必要なのは、すべての身長グループ内で、$ i $ペアの同じ身長ペアを組むときの場合の数である。

この値をどう計算するかを次に考えてみる。

すべての身長グループ内で、$ a $ペアの同じ身長ペアを組むときの場合の数

この値はすこし考えると、組み合わせ計算で$ O(1)やO(N) $で求められないとわかる。(ペアの数はバラバラで、どのペアから何人取ったかどうかは答えの計算に大事そう)となれば、DPで数え上げるという線を考えるべきである。

というわけでDPを考えてみる。ここまでくればDPの定義自体は難しくない。 $$ \begin{aligned} &dp[i][j] := グループiまで考えた時、j個の同じ身長のペアを作ったときの場合の数 \\ &dp[0][0] = 1 \\ &dp[i + 1][j + k] += dp[i][j] \times 身長グループiでkペア作る場合の数 \\ &Ans = dp[n][a] \end{aligned} $$

遷移の係数の部分では、新たに身長グループ$ i $で$ k $ペアの同じ身長のペアを作っている。このDPの計算量はこのままだと変わらず$ O(N ^ 3) $となる。

これは、以下の畳み込みの漸化式の遷移と一致してる(はっきりとDPの遷移では$ j $と$ k $が分かれている) $$ c[i + j] += a[i] \times b[j] $$

ゆえに、このDPは畳み込みをNTTすることによって、加速して$ O(N ^ 2 \log N) $にすることができるとわかる。 実際の実装の際では、$ a $ 側の係数は$ dp[i][j] $を、$ b $ 側の係数は$ 身長グループiでkペア作る場合の数 $にすればいい。

解答の見当がついたところで、$ 身長グループiでkペア作る場合の数 $を計算する。これは以下の2つのやり方があり、どちらでも構わない。

  1. $ num[i] $のうち、$ num[i] - 2k $だけペアにならない余る人を$ _ {num[i]} C _ {num[i] - 2k} $で選ぶ。ペアになる人は、二重階乗で$ (2k - 1)!! $で求まる。これらは独立なので、積を取る。
  2. 今いる人のうち、2人を選んでペアを作るをk回繰り返す。$ _ {num[i]} C _ {2} \times _ {num[i] - 2} C _ {2} \times \cdots $

いずれも適切な前計算(各身長グループで合わせて$ O(n) $)を施すことによって、畳み込みの係数を求められる。

畳み込みの順序(マージテク)

これで計算量は$ O(N ^ 2 \log N) $まで落ちた。しかし、これでもなお足りないのである。どこで計算量が使われているのかをよく見てみよう。

  • NTT1回ごとに$ O(a \log a) $　なお、$ a $はNTT時の次数の和。($ O(N \log N) $よりも厳しい評価でこう書ける)
  • $ O(N) $個の身長グループを全部NTTしてみる。これは絶対に削れない。

ここで注目すべきは、NTTの計算量はおおざっぱに$ O(N \log N) $で抑えるのではなく、その時の次数の和$ a $を用いて$ O(a \log a) $で評価していることである。 この評価から直感的にわかることは、次数が小さい順に畳み込みをすれば、余り次数を爆発させずに計算できるのではないかということである。

これはより形式的に定めると、以下の通りとなる。

  1. すべての畳み込みすべき多項式を集合$ S $に入れる。
  2. $ S $のうち、次数が低い2つの多項式を取り出し、NTTの計算を施して、計算結果を$ S $にまた入れる。
  3. これを$ |S|=1 $まで繰り返す

このテクニックは、複数個の多項式を畳み込む時、データ構造の統合の計算順序をハフマン符号のように短い順に計算するものであり、マージテクと呼ばれたりするものである。

マージテクの計算量解析はこちら。

このマージテクより、全体の計算量は$ O(N (\log N) ^ 2) $まで落とすことが可能である。

これで、この問題を時間内に計算することができた。

最後に

これでdpを高速に計算できたので、包除原理の条件ごとの値も計算できる。これが分かれば、 $$ 同じペアを0個以上含む - 1個以上含む + 2個以上含む - \cdots $$ と計算することで最終的な答えを計算することが可能である。

最後に注意すべき点として、すべての人が同じ身長で組むと、残った無制限で組める人は0人となる。この時、$ (0 - 1)!! $を計算することになるが、うまくこの例外を処理できなければWAになる(7敗) 処理は簡単で、選択の余地がないので1通りとすればよい。

ACしたソースコード

自前の畳み込みやmodint、数え上げライブラリを使用　ACLより遅い(自前のを直せ) C++

最後に

数え上げ特訓第2回。オレンジ中位の問題だけあって、1つの漸化式を立てるだけでは終わらず、より高速化できるものを見極めて適切に畳み込みやマージテクを使わなければ解けない問題。 包除の適切な条件定義も必要で、まさにDPの数え上げのテクニックのデパートであり、非常に教育的な良問だと思う。

あとけんちょんさんより詳しくかけたと思います！！！