AtCoderのビット全探索をABC167 Cで学ぶ:部分集合を漏れなく試すPython実装

読了 約15分 たびすけ
ABC167 Cで000から111までの部分集合をビット全探索する考え方を示す図

次に読む記事

関連するテーマの記事を、先に確認できます。

N冊の本を買う・買わないの二択で選ぶ問題では、候補は 2N 通りです。Nが小さいという条件なら、0以上 2N - 1 以下の整数をマスクとして使い、すべての部分集合を一度ずつ調べられます。AtCoderの ABC167 C – Skill Up では N ≤ 12 なので候補は最大 212 = 4096 個です。各maskで選んだ本の費用と理解度を集計し、すべての理解度が目標以上になる候補の最小費用を求めます。

この方法は、二択が独立して並び、候補数と1候補の判定量が制約内に収まるときに使います。候補数は N が1増えるたびに2倍になるため、任意の大きさの N にそのまま適用する方法ではありません。ABC167 Cについては、公式解説が示す全組合せの確認を、maskとPythonコードの対応まで具体化します。

N個の二択が2N通りになる理由

1冊目を買うか買わないかで2通り、2冊目でも2通りです。N冊すべてで選択すると、候補数は 2 × 2 × … × 2(N個の2)になり、2N 通りになります。たとえば3冊なら 23 = 8 通りです。空集合、1冊だけの集合、複数冊の集合、全冊を含む数です。

ここで必要なのは、候補を思いつく順に並べることではありません。各候補を重複なく数え、同じ判定を適用することです。ABC167 Cの公式解説は、買う・買わないの全 2N 通りを調べ、1候補を O(N × M) で判定する方針を示しています。そのため掲載コード全体は O(2N × N × M) になります。詳しくは ABC167の公式解説PDF の「C: Skill Up」を参照できます。

maskと本の対応を先に固定する

整数を二進数で表し、右端から数えたビットを本へ対応させます。bit 0を本1、bit 1を本2、というように、bit i を本 i + 1 と決めます。

ビット位置対応する本ビットが1の意味
bit 0本1本1を選ぶ
bit 1本2本2を選ぶ
bit ii + 1i + 1 を選ぶ

選択の判定は (mask >> i) & 1 です。maskを右へ i ビットずらし、最下位ビットを取り出しています。結果が1なら本 i + 1 を選び、0なら選びません。mask = 0 はすべてのビットが0なので空集合、mask = (1 << N) - 1 は下位Nビットがすべて1なので全選択です。

0から 2N - 1 までの整数は、それぞれ異なるN桁の二進表現を持ちます。各桁を「その本を選ぶかどうか」に対応させると、異なるmaskは異なる部分集合になります。また、どの部分集合にも選択・非選択を並べた二進表現が一つだけあるため、すべての部分集合をちょうど一度ずつ列挙できます。

3冊の構成例で8個のmaskを追う

次の例はABC167 Cの公式サンプルではなく、maskと集計の対応を確かめるために構成した例です。理解度は2種類、目標は X = 5 とします。

費用理解度1の増分理解度2の増分
本1450
本2605
本3733

表のビットは左から b3 b2 b1 と書き、右端の b1 を本1、中央を本2、左端を本3に対応させます。たとえば 011 は本1と本2を選ぶ mask = 3 です。

b3 b2 b1mask選択した本費用理解度合計達成その時点の最小費用
0000なし0[0, 0]×
0011本14[5, 0]×
0102本26[0, 5]×
0113本1、本210[5, 5]10
1004本37[3, 3]×10
1015本1、本311[8, 3]×10
1106本2、本313[3, 8]×10
1117本1、本2、本317[8, 8]10

mask = 3 では本1と本2の増分が合わさり、2種類とも5に届くので費用10を記録します。最後の mask = 7 も条件を満たしますが、費用17なので最小値は更新されません。この8行には、空集合から全選択までのすべての候補が一度ずつ現れています。

公式問題ABC167 Cへ置き換える

ABC167 C – Skill Up は、N冊の本から購入する集合を選ぶ問題です。本 i には費用 Ci があり、購入するとM種類の理解度 jAi,j ずつ増えます。購入した本による各理解度の合計をすべて X 以上にできる集合のうち、費用が最小のものを出力します。条件を満たす集合がなければ -1 を出力します。

項目公式の条件
入力すべて整数
N、M1 ≤ N, M ≤ 12
X1 ≤ X ≤ 105
本の費用 Ci1 ≤ Ci ≤ 105
理解度増分 Ai,j0 ≤ Ai,j ≤ 105
制限時間2 sec
メモリ制限1024 MiB

入力は最初に N M X、続くN行に各本の Ci とM個の Ai,j を与えます。出力は、条件を満たす購入集合の最小費用、または達成不能を示す -1 です。各本の買う・買わないをmaskの1ビットへ対応させれば、構成例と同じ列挙をこの問題へ使えます。

公式サンプル1

入力

3 3 10
60 2 2 4
70 8 7 9
50 2 3 9

出力

120

公式サンプル2

入力

3 3 10
100 3 1 4
100 1 5 9
100 2 6 5

出力

-1

公式サンプル3

入力

8 5 22
100 3 7 5 3 1
164 4 5 2 7 8
334 7 2 7 2 9
234 4 7 2 8 2
541 5 4 3 3 6
235 4 8 6 9 7
394 3 6 1 6 2
872 8 4 3 7 2

出力

1067

上の3組は、ABC167 Cの公式問題ページに掲載されているサンプルです。構成例の8行と公式サンプルは役割が異なります。前者はmaskの対応を追うため、後者は公式の入出力形式と答えを確認するために使います。

全maskで費用と理解度を集計するPythonコード

入力した本は books に費用と増分の組として保存します。外側の for mask in range(1 << n) が0から全選択までを順に作り、maskごとに費用 total_cost とM種類の合計 skill_totals を0へ戻します。内側で選択ビットが1の本だけを加算し、すべての理解度が x 以上なら費用の最小値を更新します。途中で候補を捨てる枝刈りはせず、全maskを評価します。

import sys


def main():
    values = iter(map(int, sys.stdin.buffer.read().split()))
    n = next(values)
    m = next(values)
    x = next(values)

    books = []
    for _ in range(n):
        cost = next(values)
        gains = [next(values) for _ in range(m)]
        books.append((cost, gains))

    inf = 10**18
    answer = inf

    for mask in range(1 << n):
        total_cost = 0
        skill_totals = [0] * m

        for i, (cost, gains) in enumerate(books):
            if (mask >> i) & 1:
                total_cost += cost
                for j in range(m):
                    skill_totals[j] += gains[j]

        if all(level >= x for level in skill_totals):
            answer = min(answer, total_cost)

    print(-1 if answer == inf else answer)


if __name__ == "__main__":
    main()

range(1 << n) の終端は含まれないので、maskは0以上 (1 << n) - 1 以下になります。各本の入力行から読み取った増分は、同じ本を選んだときだけ対応する理解度へ加わります。最後の all がM種類すべてを確認するため、1種類でも目標未達の候補は答えへ入りません。

コードが正しい理由を不変条件で追う

正しさは、maskの列挙が漏れないこと、各maskの集計が選択集合と一致すること、条件を満たす費用だけを最小化することの4段階で説明できます。

  1. 全候補の対応。 外側のループは mask = 0 から 2N - 1 までを一度ずつ処理します。Nビットの0/1列とN冊の選択有無は一対一に対応するため、すべての購入集合をちょうど一度ずつ調べます。
  2. 集計の不変条件。 内側の本ループで先頭からk冊を処理した直後、total_cost はそのk冊のうちmaskで選ばれた本の費用合計であり、skill_totals[j] は同じ選択本による理解度jの増分合計です。最初はどちらも0です。次の本が未選択なら値を変えず、選択されていれば費用とM項目を加えるので、この関係が保たれます。
  3. 候補の判定。 N冊を処理し終えたとき、2つの集計値はそのmaskが表す購入集合全体の合計です。all(level >= x for level in skill_totals) が真なら、問題の条件を満たす集合として費用を候補にします。
  4. 最小値と達成不能。 条件を満たした候補だけで answer の最小値を更新するので、全候補の処理後には達成可能な購入集合の最小費用が残ります。一度も更新されなければ達成可能な集合がなく、-1 を出力します。

計算量と、空集合・全選択を含む境界

maskは 2N 個あります。各maskで最大N冊を確認し、選択した本について最大M種類を加算するため、掲載コードの時間計算量は O(2N × N × M) です。公式制約の N ≤ 12 ではmaskの最大数は4096個です。これは候補数の算術であり、Pythonの実測時間や提出結果を保証する値ではありません。

books はN冊分の費用とM種類の増分を保持するので O(N × M) のメモリを使います。各maskで作る skill_totalsO(M)、その他の変数は O(1) です。全体では O(N × M) と表せます。

  • 空集合:mask = 0 も評価します。公式制約では X ≥ 1 なので合計理解度がすべて0の空集合は条件を満たしませんが、列挙から除外しません。
  • 全選択:mask = (1 << N) - 1 も評価します。すべての本が必要なケースでも、最後の候補として判定できます。
  • 達成不能: どのmaskでも全M種類が X 以上にならなければ answer は初期値のままで、出力は -1 です。
  • 最小サイズ:N = 1M = 1 でもmaskは0と1の2候補で、同じコードを使えます。
  • 増分0:Ai,j = 0 は公式制約内の値です。その本を選んでも該当する理解度は増えないため、0を未入力や未選択とみなしてはいけません。
  • 公式制約外の参考境界:X = 0 はABC167 Cの公式制約外です。この値を導出テストとして入れると、空集合の合計0が条件を満たすため費用0になります。これはコードの境界を確認するための挙動であり、公式仕様の入力例ではありません。

次の練習:ビットの対象を切れ目へ変えるABC197 C

ABC167 Cでbitに対応させたのは「本を選ぶかどうか」でした。次の ABC197 C – ORXOR では、長さNの数列を一つ以上の空でない連続区間へ分け、区間ごとのbitwise ORをすべてbitwise XORした値の最小値を求めます。公式制約は 1 ≤ N ≤ 200 ≤ Ai < 230 です。

この問題で選ぶ対象は要素そのものではなく、隣り合う要素の間にあるN-1個の切れ目です。それぞれに「仕切りを置く・置かない」の二択があるので、maskは 2N-1 個です。切れ目を決めると最後の区間は自動的に確定するため、ビット数はNではなくN-1になります。ABC167 Cの「二択列を整数へ対応させる」考えを、対応先だけ変えて使う練習です。

公式サンプル1で確認すること

ABC197 Cの公式問題 のサンプル1は、入力が N = 3、数列が 1 5 7、出力が 2 です。公式解説の 分割全探索の説明 を読み、コードを書く前に「なぜ仕切りの候補がN-1個なのか」「選んだ仕切りの後に最後の区間をどう確定するのか」を説明できれば、maskの意味を固定する段階まで進めます。ABC197 Cの完全コードや、ここで扱っていない演算の体系的な解説はこの橋渡しの範囲に含めません。