Codeforces Round #621 (Div. 1 + Div. 2) E. Cow and Treats
問題概要
本の草が横一列に並んでいて、 本目の草の味は整数 で表現される。また 頭の牛がいて、 頭目の牛は味 の草を好み、それを 本食べると満足する。
以下の処理を行う。
- 互いに素な牛の集合 を選ぶ。一方または両方が空であってもよい。
- の牛は左端から、 の牛は右端から、好きな順番で出発する。
- 各牛 は訪れた草のうち味 のものだけを食べながら進み、 本の草を食べると満足してその場所に留まる。草は一度食べられると消滅する。別の牛が留まっている場所は通過できない。
出発順を適切に選んだ時に に含まれる全ての牛が満足する時、この を良い集合対と呼ぶことにする。良い集合対における の最大値を求め、さらにその最大値を達成する良い集合対の総数を で求めよ。
制約
- 同一の を持つ牛は2頭以上存在しない
解法
境界の全探索と重複の排除
まず、同じ味を好む牛は左右それぞれに高々1頭しかいません。2頭以上いる時には、必ず最初に出た牛に他の牛が通せんぼされるからです。
左側の先頭の牛が進む位置と右側の牛が進む位置の境界を 通り全探索することを考えます。
この境界を決めると、境界の左右それぞれに存在する各味の草の本数が決まります。そして各味ごとに、 がその本数以下であるような牛を左右それぞれ1頭または0頭採用することになります。逆に全ての味についてこれを満たす選び方をしていれば、止まる場所が遠い順に牛を出発させることで、必ず進路を妨害されることなく全ての牛を満足させることができます。
つまり集合に含める牛は味ごとに独立に決められるので、味ごとに採用する牛の頭数(0~2頭)の最大値を求めて全て足せば、その境界に対応する牛の頭数の最大値が求められます。これを全ての境界について試せば良さそうです。
ただし今回の問題では場合の数も要求されているため、重複を排除する必要があります。例えば上の図のような状況では、赤い境界線が図より1つ右側にあるときでも同じ場合が数えられてしまいます。
これを防ぐために、
- 境界が左端にある場合を考える時は、左側から出発する牛は存在しない。
- それ以外、つまり境界の左隣に草がある時は、左側の先頭の牛がちょうどその草を食べて止まるような場合のみを数えることにする。
というルールを追加します。このルールによって、例えば上の図のような牛の選び方は図に示した位置の境界を考えている時だけカウントされることになり、重複を排除できます。
具体的な処理アルゴリズム
具体的な処理の仕方を考えましょう。まず各牛を味ごとに分類し、同じ味の中では が小さい順にソートしておきます。制約より、ここで味も も等しい牛が複数存在することはありません。
左右の境界を全探索します。そして各味ごとに、左右それぞれの領域に何本の草があるかを累積和や差分更新などで求めて、採用する牛の頭数の最大値とその選び方の数を以下のように求めます。
見ている味が、境界の左隣の草の味ではない時
各味について、その味を好む牛はあらかじめ の値でソートしておきました。これを左右それぞれに存在する草の数で二分探索することで、小さい方から何頭目までの牛が候補になるかを計算できます。この頭数を左右それぞれ としましょう。
このとき、左右ともに1頭ずつの牛を採用する場合の数は と計算できます。左右それぞれ独立に牛を選んだペアの個数から、同じ牛を選んでいるものを除くという計算です。この値が正である時は牛を2頭採用できます。
2頭採用できないならば、1頭だけを選ぶ場合の数は で、これが正なら1頭採用になります。これも ならば0頭採用です。
見ている味が、境界の左隣の草の味である時
この時は左側の牛は完全に固定されます。境界の左側に今見ている味の草が 本あるとすると、ちょうど の牛を採用する必要があります。これがもし存在しないならば、この境界で考えるべき場合の数はそもそも存在しないので、処理を中断します。
その上で右側の牛として採用できる牛の候補を二分探索で求めます。もしこの候補の中に先ほど左側で採用した の牛が存在する場合は取り除く必要があります。
右側の牛として採用できる候補の数を として、これが正なら2頭採用で場合の数は 、 なら1頭採用で場合の数は になります。
これで各味についての牛の頭数の最大値と場合の数が求められたので、それぞれ総和/総積を取ることで全体での最大値と場合の数を求められます。あとは全ての境界についての結果をマージすれば答えを求めることができます。