何の話かと言うと
この問題、公式解説では、最小カット問題(=最大フロー問題)に帰着して、Dinic 法で解くという解法が紹介されていますが、もうちょっと素朴な「ナップサック問題」としても解ける(解けた)ので紹介します。
ナップサック問題との類似性
まず、n 番目の家には、n+1 番目以降の家の鍵が置いてあることから、家を訪れる順番は、基本的に番号順になります。つまり、n=1,2,...,N について、それぞれの家に入る/入らないという判断をしていきます。
また、n 番目の家を入る/入らないという判断をするごとに、この後に入ることのできる家に対する制約条件が変化していきます。
これは、ナップサック問題に類似した状況です。
ナップサック問題では、n=1,2,...,N 個の荷物のそれぞれについて、入れる/入れないという判断をしていきますが、判断をするごとにナップサックの残り容量が減って、追加で入れられる荷物に対する制約が変化していきます。
このようなナップサック問題では、残り容量をキー、その時点での価値をバリューとするディクショナリーに情報を記録していくのが効率的です。
今回の場合は、未収集の家の鍵が徐々に減っていくので、たとえば、「現在残っている鍵の状態」をキーにすることが考えられます。最大 N 種類の鍵があるので、それぞれについての残り数を並べたタプルをキーにしてみましょう。初期状態は、次のように構成します。
import sys, copy def main(f): global N, W, A, C N, W = map(int, f.readline().split()) A = [None] + list(map(int, f.readline().split())) C = [[] for _ in range(N+1)] num_keys = [0] * N for n in range(1, N+1): C[n] = list(map(int, f.readline().split())) C[n] = C[n][1:] for c in C[n]: num_keys[c-1] += 1 num_keys = tuple(num_keys) dp = defaultdict(lambda: -10**30) dp[num_keys] = 0 print(dp) #main(sys.stdin)
# input 5 5 5 2 10 3 6 1 3 1 3 0 1 5 0 # output {(0, 0, 2, 0, 1): 0}
この例では、3番目の家の鍵が2本、5番目の家の鍵が1本残っており、現在の価値()は 0 です。
この後は、n=1,2,...,N のループを回しながら、ディクショナリー dp のキーを取り出して、その状態から n 番目の家に入る場合、入らない場合の状態を追加していきます。家に入った場合は、その家にある鍵を減らしていきます。もちろん、残りの鍵が0でない家には入れません。
で、これを(効率はあまり考えずに)素朴に実装して提出するとこうなります。
AC x 10, TLE x 2 で、予想通り(?)TLE していますが、max ケースでも AC しているものもあり、方向性としては悪くなさそうです。
状態数の削減
この手のナップサック問題の場合、n が増えるごとに状態が倍になるので、状態数が爆発的に増えることにより、状態(ディクショナリーのキー)に関するループが長くなります。なんらかの方法で、状態数を減らす必要があります。
例えば、普通のナップサック問題であれば、ナップサックの総容量 W とすると、状態は 0 〜 W の整数に限定されるので、W が小さければ、自然と状態数も少なくなります。(原理的には、状態数は倍々で増えますが、結果的に残り容量の値が被る場合が多発します。)
では、今回の場合、何らかの方法で「状態」の範囲を限定して、状態数を減らすことができないでしょうか?
まず、「鍵の残り数」で考えることをやめます。重要なのは、残り数が0か0でないか(ある家に入れるか入れないか)、という事で、0 でない場合の具体的な値は本質ではありません。
・・・でも、残り数をカウントしないと、途中経過がトラッキングできないのでは・・・・?
そうでもありません。たとえば、家1に家3の鍵がある場合、家1に入らなければ、その時点で、家3に入れないことは確定します。(家3の鍵の残り数がいくらかは考える必要がありません。)つまり、それぞれの家に「入れる/入れない」という0/1のフラグをN個持っておき、ある家に入らないという選択をするごとに、入らなかった家にある鍵の家のフラグを1にすればよいのです。こうすれば、 の整数値をディクショナリーのキーにすることができて、状態数を減らすと共に、(重い)タプルの操作を無くすことができます。
・・・いやそれでも、 ってめちゃめちゃ大きくないです?
確かに・・・・。
いや、まだ諦める必要はありません。ここがポイントなのですが、今回の場合、それぞれの家にどの鍵があるかは(実際に入って確かめなくても)事前にすべてわかっています。「家1に家3の鍵がある場合、家1に入らなければ、その時点で、家3に入れないことは確定します」と言いましたが、家3に入れないということは、家3に鍵がある家にも入れない、という事実が確定します。つまり、1つの家に入らないと、その先で入れなくなる家が連鎖的に増えていきます。こう考えると、状態フラグは必然的に1が多くなり、実際に発生する状態の数は、 よりもずっと小さくなる可能性があります。
さらに、n 番目の家の処理がおわれば、この後は、n+1 〜 N 番目のフラグのみが必要で、それ以前のフラグの値は参照する必要がありません。そこで、n 番目を家の処理を終えたら、n 番目のフラグは強制的に 0 にします。こうすれば、処理がすすむごとに状態フラグの(実質的な)桁数が減っていき、状態数もどんどん減っていきます。
このあたりの処理をまとめて実装するとこんな感じです。連鎖的に入れなくなる家は、事前に計算して、bitmask として用意しておきます。
mask = [0] * (N+1) # mask[n] : 家 n に入らなかった時に(連鎖的に)入れなくなる家を示す bitmask for c in C[N-1]: mask[N-1] |= 1<<(c-1) for n in range(N-2, 0, -1): for c in C[n]: mask[n] |= 1<<(c-1) mask[n] |= mask[c] st = 0 dp = defaultdict(lambda: -10**30) dp[st] = 0 for n in range(1, N+1): dp_n = defaultdict(lambda: -10**30) for st in dp.keys(): # 家 n に入らない場合 st_n = st st_n |= mask[n] # 連鎖的に入れなくなる家のフラグを立てる st_n ^ 1<<(n-1) # n 番目のフラグを強制的に 0 にする dp_n[st_n] = max(dp_n[st_n], dp[st]) if st & 1<<(n-1) == 0: # 家 n に入る場合 r = A[n] - W dp_n[st] = max(dp_n[st], dp[st] + r) dp = dp_n
で、これを提出すると・・・
惜しいです。TLE が 1 つだけ残ります。
状態の刈り込み
最後の手段で、状態の刈り込みを行いましょう。
ある2つの状態 A, B を比べて、
・A の方が B よりも状況的に不利で、かつ、その時点の価値も B より低い
という状況があった場合、状態 A をそれ以上トラッキングする必要はありません。その状態はディクショナリーから削除できます。
今の場合、
・B で入れない家は、すべて、A でも入れない
という場合、A は B よりも状況的に不利です。ビット表現であれば、flag_A | flag_B == flag_A と表せますね。このような組 A, B について、さらに dp[A] <= dp[B] であれば、A を捨てることができます。
ただ、今回の場合、このチェックを漏れなく行うには、状態数を として、 のループが必要になります。 が大きいとこの処理自体が遅くなる可能性もあります。
で、ここはぶっちゃけ状況しだいなのですが、仮に、この刈り込みが初期段階から強く効いたとすると、刈り込みによって状態数 が抑えられて、刈り込みの処理自体も遅くならずに済む可能性もあります。とにかくやってみましょう。
無事にすべて通りました。