🤪

atcoder ZONE2021 C問題をRubyで解く

2021/05/03に公開

Ruby

atcoder ZONE2021 C問題 "MAD TEAM"が解けなくて超悔しかったので、忘れる前に記事化しておきます。
以下、「羅列」という単語が頻出します。寝ぼけてただけなので「列挙」と読み替えてください。

大まかな解法

まず思いつくのは愚直法。普通に候補者の中からチームメンバーを3人選び出し、総合力を比較するだけです。
候補者は最大3000人なので、三人を選ぶ組み合わせは(3000×2999×2998)/(1×2×3)≒(3000×3000×3000)/6=3000×3000×500=4500000000=4.5e9通りくらいあります。atcoderの実行制限時間で羅列できるのは7桁くらい(rubyなら6桁)までなので、10桁は間に合うわけないですね。

正解の解き方は2つあり、一つは「二人を決めると最適な最後の一人は自動的に決まる」。これは(3000×2999)/(1×2)≒4.5e6通りで決まりますが、rubyでは実行制限時間に間に合うか不明なので解説しません。というか最後の一人を決める方法がよくわからなかったので解説できません。
もう一つは能力値の方を二部探索。取り得る能力値の組み合わせは1e9通りなので二部探索しても30回くらいで終了します。(毎回候補者を全員見て、(30×29×28)/(1×2×3)≒1e3通りくらいの組み合わせを羅列する必要がありますが)

二部探索って何だよ

データが順番通りに並んでおり「ここから左は全部false、右は全部true」(またはその逆)みたいな場合に使う方法です。雑に言うと「データの真ん中を調べる→正解が左にいるか右にいるかがわかる→正解がある方の半分のデータのうち、真ん中を調べる」の繰り返しです。詳しいことはググってください。
rubyでは配列クラスにbsearchメソッドとして実装されており、実際これを使ってACしている人もいます。私も使ってみましたが何故かTLEしたので解説しません。

今回の場合、能力値0の人はいないため総合力1(以上)のチームは絶対に作れます。能力値1e9+1の人はいないため総合力1e9+1のチームは絶対に作れません。つまり
「作ったチームで総合力1のボーダーはクリアできる」
「作ったチームで総合力2のボーダーはクリアできる」
……
「作ったチームで総合力X-1のボーダーはクリアできる」
「作ったチームで総合力Xのボーダーはクリアできる」
「作ったチームで総合力X+1のボーダーはクリアできない」
……
「作ったチームで総合力1e9のボーダーはクリアできない」
「作ったチームで総合力1e9+1のボーダーはクリアできない」
となっています。
このXの値を二部探索で探します。

で、作ったチームの総合力がボーダーをクリアしてるかってどう調べるのよ

一番わかりやすいのは、候補者一人で5能力全てクリアしてくれる人がいることです。[100,100,100,100,100]がいてくれれば、残りの二人がカスでも総合力は100ですね。
しかしそう上手くはいかないので、二人や三人の組み合わせを考える必要が出てきます。いちいち候補者3000人から3人選抜して調べてたら意味ないです。
そこで、何故組み合わせる必要があるのか考えます。総合力2のボーダーをクリアしたくて、メンバーの二人が[100,100,100,1,1][100,100,1,100,1]だった場合、最後の一人に[100,100,100,100,1]と[1,1,1,1,2]どちらを選ぶべきでしょう？

今回のチームメンバーは似たような人材を固めて一芸特化するのではなく、互いに短所を補完しあう必要があります。上の例で[100,100,100,100,1]を選んでしまった場合、三人とも発想力(5つ目の能力)が短所になってしまいます。[1,1,1,1,2]の人は発想力が長所なので、こちらを選ぶのが正解ですね。

「長所・短所って何だよ。[100,100,100,100,3]の人は発想力が短所だけど[1,1,1,1,2]より発想力高いぞ」という疑問にもお答えします。今回は「クリアしたいボーダーを超えている能力が長所、超えていない能力が短所」です。先程の例ではクリアしたいボーダーは2。つまり[100,100,100,100,3]の人は5能力全て長所です。仮にボーダーを4に引き上げたら[100,100,100,100,3]の人は発想力だけ短所で他が長所、[1,1,1,1,2]の人は5能力全て短所です。

このようにカッチリ定義すると、全ての能力は長所短所の二通りしかありません。能力は5つなので「オール短所」「発想力だけ長所」……「オール長所」の2の5乗で32通りにカテゴライズできます。
私はここまで調べても「それの何が嬉しいんだ？　候補者が3000人いたら選び出す手間は一緒だろ」とか思っていました。先程「何故組み合わせる必要があるのか」に対し、なぜなら「互いに短所を補完しあう必要」があるからだと答えました。逆に言うと、「全く同じ長所短所を持った候補者を二人も三人も入れても意味がなく、全く同じ長所短所を持った候補者は一人いればいい」わけです。
つまり32カテゴリに分類された3000人全員から三人選ぶ組み合わせを考える必要はなく、「この長所短所の組み合わせの候補者はいる、この組み合わせの候補者はいない」で考えれば良いわけです。

やっと本題にたどり着きました。今回の大まかな解法は以下のようになります。

ボーダーを二部探索で決める(大ループ)
そのボーダーを元に候補者の長所短所を見て32通りにカテゴライズ(小ループ)
「この組み合わせカテゴリの候補者はいるよ！」と記録
2に戻る。候補者の人数分なので最大で3000回くらい
「どの組み合わせカテゴリの人材がいるか、いないか」のリストが完成
三人以下でチームを羅列する。最大で32通りから3人なので1000通りか2000通りくらい
5能力とも長所のチームがあるか調べる
あればボーダーを上げて1に戻る。なければボーダーを下げて1に戻る。1から1e9+1までで二部探索するので多分30回くらい

bit演算で、チームが5能力とも長所になっているか高速判定

5能力がそれぞれ長所か短所か、rubyでは長さ5の配列にbooleanで入れてもいいと思います。Array#zipやArray#ary?を使えば実装も直感的でしょう。
しかし私はTLEが怖いのでbit演算を使いました。多分こんな事しなくてもACだと思います。試していませんが。

偶奇性を調べたり、2のべき乗か調べたりするのに便利なbit演算ですが、今回はカテゴリ分けに使いました。パワー・スピード・テクニック・知識・発想力をそれぞれ1の位、2の位、4の位、8の位、16の位に振り分け、長所なら1短所なら0で表現します。5能力が全て長所なら2進数で11111(十進法で31)、全て短所なら00000(十進法で0)、長所長所長所長所短所だった場合は2進数で01111(十進法で15)、短所短所短所短所長所なら2進数で10000(十進法で16)になります。左右が逆でややこしいですね。
で、これをor演算します。2進数で01111のメンバーと10000のメンバーを組み合わせれば11111でボーダークリアです。「足し算じゃ駄目なの？　01111+10000=11111でボーダークリアじゃん」と言われそうですが駄目です。00001のメンバーと00011のメンバーを組み合わせると能力の有無は00011のはずですが、加算すると繰り上がって00100のチームになってしまいます。テクニックの能力どこから来た。

具体的なコード

モジュール解説は後に回すとして、48~49行目で入力を受け付けて変数に入れています。候補者の定数名をMEMBERとしているせいで候補者なのかチームメンバーなのかややこしすぎる。皆さんはcandidateとteammateをきちんと使い分けてください。

ボーダーを二部探索で決める(大ループ)

55~56行目が二部探索のループです。前述の通り「総合力Xのボーダーはクリアできる」「総合力X+1のボーダーはクリアできない」となったら終了です。
最初「初期範囲もっと絞り込んだ方がいいだろ」と思い、入力配列から最小値と最大値を算出していました。しかし1ループで範囲を半分に絞れるのが二部探索の美点。逆に言うと下ごしらえして範囲を半分まで絞り込んでおいても1ループ分しか削減できません。配列の最小値と最大値を計算する分むしろ遅くなってしまいました。

そのボーダーを元に候補者の長所短所を見て32通りにカテゴライズ(小ループ)

61行目です。メンバーにボーダー(変数名mid)を送り、能力カテゴリ0~31までのうちどれかを申告させます。カテゴライズするメソッドを定義しているのは7行目。ただでさえTLEにビビってbit演算を使っていますが、更にwhileループを使って高速化しています。eachがどうこう書いてますが、mapでbool値の配列にしても良かったんじゃないかこれ。

「この組み合わせカテゴリの候補者はいるよ！」と記録

0から31までの数字にカテゴリ分けしたからと言って、それを素のまま長さ3000の配列にしたら意味がありません。必要なのはいるかいないかだけなので、「全カテゴリの人材がいないよ」の長さ32の配列を予め57行目で作っておき、「このカテゴリの候補者がいたよ」という場合を63行目でtrueにしていきます。変数名がcountなのは迷走していた頃の名残り。いるかいないかなので皆さんはexistを使ってください。62行目のコメントが微妙にわかりにくいのもメンテナンス忘れです。
65行目のガード節は一人で全能力クリアしている人が見つかった場合です。このボーダーはクリアできるに決まっているので、それ以降を無視します。

2に戻る。候補者の人数分なので最大で3000回くらい

66行目です。

「どの組み合わせカテゴリの人材がいるか、いないか」のリストが完成

ループが終わると、57行目で作った配列がいい感じにtrue,falseで構成されています。

三人以下でチームを羅列する。最大で32通りから3人なので1000通りか2000通りくらい

羅列する前に、全能力がボーダー未達の人はいてもいなくてもチーム総合力には無関係、いるだけ組み合わせが増えて羅列に時間がかかるので66行目で消します。
22行目のメソッドで羅列します。まず候補者が0人(全員が全能力ボーダー未達で、66行目で消しちゃった)ならクリアできるわけないので早期リターン。また、一人で全能力クリアしている候補者がいる場合はクリアしているに決まっているので早期リターン。
rubyでは組み合わせの羅列はArray#combinationが使えますが、[false,true,false,true,....]の形ではcombinationが使いにくいので29~30行目で[1,3,...]の形に作り直します。
また、Array#combinationは長さ2の配列から3つを選ぼうとすると空配列を返すので、長さ3以上の場合だけ3つ選んで、長さ2以下の時は2つ選んでねと言っています。これだと長さ1の時は空配列を返してしまいますが、一人で全能力クリアできる場合は早期リターン済みなので空配列を返してfalseになっても平気です。「長さが2ならcombination使わなくても、要素2つとも選んでor演算すれば良くない？」と思いますが、なぜかそっちのほうが遅かったので解説しません。

5能力とも長所のチームがあるか調べる

34行目と36行目です。
人材がいるカテゴリを2つまたは3つ選んでチームを編成。各カテゴリをor演算。十進法で31、2進数で11111になっているチームがあればtrue、なければfalseです。

あればボーダーを上げて1に戻る。なければボーダーを下げて1に戻る。1から1e9+1までで二部探索するので多分30回くらい

70行目以降。
最終的に「総合力Xのボーダーはクリアできる」「総合力X+1のボーダーはクリアできない」が判明したらXを表示して終了。

後書き

実行時間91msは過剰高速化ですね。早期リターンを消したり、ブロックを使って分かりやすくしても大丈夫でしょう。(分かりやすいバージョンを書かないコーダーの屑)
あとc++もpythonも読めない僕のために、rubyistはもっと解説記事を書いてください。(私利私欲)