🤪

Atcoder EducationalDPContest O MatchingをRubyで解く

2021/05/22に公開

Ruby

AtCoderの実行時間はC++基準で決められているため、遅い言語で書くと想定解法通りでも間に合わない事があります。
この問題も(多分)その内の一つで、ACするには(多分)やや難解な書き方が必要です。
私は先人のRubyのコードを頑張って解読したので、忘れないうちに記録します。

問題

想定解法

まずはbitDP

https://algo-logic.info/bit-dp/
今回の問題の入力例1に当てはめてみます。二次元配列にして行が男で列が女。
男側は「男0」「男1」「男2」で、三人なら三行。一方で女側は女0女1……ではなく「女は一人もマッチングしていない000」「女0のみがマッチング済みの001」「女1のみがマッチング済みの010」「女0と女1がマッチング済みの011」……「全員マッチング済の111」で、三人なら2の三乗で八列です。
まず男1がいる0行目(以下、x行目とかx列目は全部0始まりで数えます)。男1と相性がいいのは女2と女3なので010列目(2列目)と100列目(4列目)を1にします。
次に1行目。男2と相性がいいのは女1と女3です。男1と女2、男2と女1がマッチングした場合の011列目(3列目)、男1と女2、男2と女3がマッチングした場合の110列目(6列目)、男1と女3、男2と女1がマッチングした場合の101列目(5列目)を1にします。
最後に男3。これは最後に残った女が誰でもマッチが成立するので、1行目の3パターンを全て合計してOKです。女は123全員がマッチ成立済みなので、111列目(7列目)を三回インクリメントして答えは3です。
このように列側について一人一列ではなく、集合の状態ごとに一列を与えるのがbitDP……なんでしょうか？　詳しい人がいらっしゃいましたら、このクソザコ緑コーダーにごコメント下さいませ。

高速化の工夫

bitDPベタ書きだと、C++やPythonでもTLEになるらしいですね。
よく考えたら男が一人目の時は女も一人目(マッチング済みの女はいない)に決まっていますから、101列目とか111列目とか考える必要はないわけです。
同様に男が二人目の時は女も二人目、マッチング済みの女は一人に決まっているわけで、一つ上の行で見る列は001か010か100、インクリメントする列は011か101か110の三通りしかないですね。
このように、二次元配列には明らかに男女の人数が合っていないマスが結構あるわけです。人数が合っていないマスを踏んだ時に、さっさと次の列に行くようにするとかなり高速化出来ます。上に挙げた解説ではhamayanさんのサンプルコードがこの方法です。
さらに一歩進めると、「そもそもこのDP、行は必要か？」という考えにたどり着きます。101列目になれば男1はとっくにマッチング済に決まっているので、いまさら0行目をいじったりしません。逆に001列目では男3の出番はまだ先、つまり2行目を触る事もありません。「男がi人目ということは男はi-1人目までマッチング済み、つまり女もi-1人目までマッチング済みの列を探してそこに一人加えてi人目にする」というアルゴリズムにはそれだけ無駄が多いわけです。
そこで男と女を逆転し、「女がi人マッチする状態という事は男もi人目に決まっている」と考えます。すると列を見るだけでマッチング済の男が何人、これからマッチする男が何人目かが一意に決まる事に気づきます。なので行をなくして、女のマッチング状況だけの一次元配列にできます。上の解説ではこちらの解き方が主流です。C++もPythonも読めないので憶測ですが。

RubyでACするために必要な工夫

で、ですよ。
Rubyでここまでの工夫を全て行っても、ACは(多分)取れません。

実際、Rubyでこの問題をACしている人は私の前には二人だけです。(私が3人目)
先人のACコードで使われているテクニックが以下です。

(逆)popcountのメモ化

Rubyにはpopcount関数がありません。スクラッチするならInteger.to_s(2).count("1")と書くのが一番早いらしいですが、これでもまだスピード不足です。
幸いAtCoderはメモリが豊富なので、配列を用意しておいて一度計算したpopcountはメモしてしまえば高速化できます。
まあ今回は1を数えるpopcountではなく、逆に0を数えてメモにするのですが……。

整数xと整数-xのand演算

ここが一番何やってるかわかりにくかったです。基礎基本がなってなかったのでゼロから勉強しました。

bit反転

2進数で表現して、0と1を入れ替えます。18を2進数で表すと10010、これをbit反転すると01101(十進法で13)です。

2の補数

2進数で10010と01101の和は11111です。この例に限らず、元の整数xとそれをbit反転した数字(Rubyでは~xと表します)を足すと、当然2進法で1の羅列になります。
ではbit反転した数に1を足した和~x+1を足すとどうなるでしょう。
18 + (~18 + 1)
2進数に直して
11010 + (~11010 + 1)
カッコの位置をずらして
(11010 + ~11010) + 1
先程の定義より
11111 + 1 = 100000
オーバーフローを無視して
00000
これがbit上で負の数を扱う方法です。bit上で-xは~x+1で表現され、これを2の補数と呼びます。

元の数と2の補数をand演算

ここからが本題です。2進数で11010のbit反転は00101、2の補数は+1して00110です。では最初の数とこの2の補数をand演算するとどうなるでしょう。
11010 & 00110 = 00010
他の場合についても考えてみます。
1101 & -1101 = 1101 & (0010 + 1) = 1101 & 0011 = 0001
101000 & -101000 = 101000 & (010111 + 1) = 101000 & 011000 = 001000
10 & -10 = 10 & (01 + 1) = 10 & 10 = 10
元の数字の中で、一番右に1が立っているビットだけを抜き出せているのがわかるでしょうか？
証明は他に譲りますが(怠慢)このテクニックを使うと1の立っているbitを高速で検索できます。

……これ、情報系の学校とかで教わるんですか？

問題に転用

今回の問題では「マッチングの見込みがある女」についてループする場面があります。
「全女に対してループを回し、マッチングの見込みがない女に関しては早期リターン」というのが想定解法ですが、Rubyでこれをやっても間に合いません。
そこで使うのが上で出てきたアンド・マイナスです。マッチ可女を1、マッチ不可または計算済みの女を0としておき、補数との論理積を取って1の立っている女を高速で抜き出し、計算し終えたら0にします。
この考えでは未マッチの女を1、既マッチの女を0と表現するため、popcountしたら未マッチの女の数が出てしまいます。「女がこういう残り方してる場合、男は何人目よ？」「というか既マッチの男は何人よ？」「つまり既マッチの女は何人よ？」という計算には、女の全数からpopcountした数を引く必要があります。ただでさえ遅いInteger.to_s(2).count("1")から更に引き算していたら絶対に間に合わないため、「全数 - popcount」の数を配列にメモしておくわけです。

ACしたコード

ほぼ写経です。

入力を受け付けて2のべき乗に

13行目から16行目。
わざわざONE_TEAMと書いているのは、ただのMEMBERS_COUNTだと男女合わせた数だと勘違いしそうだったからです。今見るとこの定数全く使ってないし、受け取ったらそのままべき乗すれば良かった。
15行目で「大して」が「対して」になってますね

動的計画法用の配列を作成

17行目から20行目。
dp[111……1]が「誰もマッチしていないパターン数」。1通りに決まっています。
dp[0]が「全女マッチ済みのパターン数」。最終的にここに答えが入ります。

逆popcountのメモ用配列作成

21行目から25行目。
111……1なら0は0個に決まっています。
「間に合うんじゃないか？」とか言ってますが、この後実験してみたら(確か)間に合いませんでした。

DP配列の末尾から先頭に向けて大ループ開始

27行目から29行目。
複数形の変数名に注目。これは女一人を表現した変数ではなく、前述した「未マッチなら1、既マッチなら0」で女全員の状態を表現しています。これが00……0、つまり女全員が既マッチになったら終了です。

未マッチ女がこのような状態の場合、次にマッチする未マッチ男は誰か

33行目。
C++のサンプルコードではpopcountでいちいち計算しています。Rubyでは間に合わないし、今回数えるのは1ではなく0というのは前述の通り。
ループ内で計算して当てはめていきます。代入されていない場合は無視してOK。

マッチ可の女を2進数で

36行目。
未マッチであろうと相性が悪ければNGです。and演算して、未マッチかつ相性のいい女だけ1、片方でもダメなら0にします。

マッチ可の女を一人ずつ抜き出して小ループ

37行目から40行目。
ここで出てくるのがさっきの補数との論理積を取るテクニックです。
一度考えた女は48行目で引いてループしています。

この男女でマッチングしたと仮定した場合、女の未マッチ・既マッチの状態はどうなるか

42行目。
単純に引くだけ。

その未マッチ・既マッチ状態のパターン数を増やし、またその状態で次にマッチする男が誰かメモする。

44行目から46行目。
DPは元のパターン数を遷移先のパターン数に足し、逆ポップカウントのメモは元の数+1の代入。

出力

52行目。
「なぜここで剰余を取っているのか？」と思うかもしれませんが、実はループ内でいちいち剰余を取っていると2秒に間に合いません。
書き忘れましたが、大ループでdowntoを使っていないのも似たような理由です。Rubyではwhileループが一番速いです。

競プロ強い人にお願いです。
後進のためにコメント書いてくださいとは言いませんから、後進のために変数名分かりやすくしてください……。

問題