ABC164 F - I hate Matrix Construction をSAT/SMTソルバで解く

AtCoder Beginner Contest #164 のF問題として出題された**「I hate Matrix Construction」は非常に難しく、AtCoder Problems による推定難易度では新体制移行後のABCで2番目に難しい問題**となっています。

その要因として、問題の考察自体が難しいことも挙げられますが、それよりも解を得るために必要な盤面のパターンや条件分岐が膨大であり、実装が非常に困難であることが大きいです。

ところで、この問題は指定された条件を満たす解を与えるプログラムを考える、「構成」系の問題です。特にこの問題の条件は非常に単純な論理式で記述することができます。このような問題はSATやSMTなどのソルバと非常に相性が良いです。

SMTソルバを用いることで、**問題に対して一切本質的な考察を行うことなく、満たしてほしい制約を入力するだけで自動的に問題を解いてくれます。**実際にやってみました。

なお、SATやSMT自体の解説は SAT/SMTソルバを自作してみる話 - るくすの日記 ~ Out_Of_Range ~ あたりがわかりやすいです。

愚直に実装する

SMTソルバZ3のPythonバインディングであるZ3Pyを使って実装してみます。

solver = Solver()
cs = [BitVec("c{}".format(i), 64) for i in range(n * n)]

for direction, (ops, results) in enumerate([[ss, us], [ts, vs]]):
    for i, (op, result) in enumerate(zip(ops, results)):
        if direction == 0:
            row = [cs[i * n + j] for j in range(n)]
        else:
            row = [cs[j * n + i] for j in range(n)]
        solver.add(reduce(and_ if op == 0 else or_, row) == result)

solver.check()

※入出力を省略しています。完全なファイルはこちら

これだけです。なんと、入出力を除けばわずか10行でこの面倒な問題を実装できてしまいました。

小さなケースで試してみましょう。

$ cat in/small.txt
5
1 1 0 1 0
1 1 1 0 0
18444487875849355213 18426372321699684173 72057740066822216 18444416407593549773 72339103374907464
18446632953678127084 18444421836365094892 18446708820012949484 4684025131465447500 4611826764505684044

$ time python abc164_f.py < in/small.txt | tee out/small.txt
0 35184908959744 18085866798841984 4684025131465455693 18354027719693426508
0 0 5497558138880 4684025131465447501 18354027719693426508
74309539880507496 72057740066822248 74309539880507464 18446744073709551615 9223372036854251519
4611826756989485060 4611861941898444804 4629841155532521604 4684025131465455692 18354027719693426509
13832554250846068680 13832559748437761992 13814539852335538024 9223372036854775807 13927915155377109247

real    0m2.169s
user    0m1.438s
sys 0m0.656s

$ crystal run abc164_f.validate.cr -- in/small.txt out/small.txt
AC

問題なく生成できてますね。Z3サイコー! (abc164_f.validate.crは私が実装したこの問題の簡易バリデータです)

高速化してみる

上のソルバーを問題の制約である $N=500$ のテストケースで試すと、メモリを食いつぶして落ちます。

ちょっとした工夫をして、メモリと時間を節約するように書き直したのが以下のものです。

ans = [0] * (n * n)

for bit in range(64):
    solver = Solver()
    cs = [BitVec("c{}".format(i), n) for i in range(n)]

    for i, (op, result) in enumerate(zip(ss, map(lambda v: (v >> bit) & 1, us))):
        if op == 0:
            if result == 0:
                solver.add(cs[i] != (1 << n) - 1)
            else:
                solver.add(cs[i] == (1 << n) - 1)
        else:
            if result == 0:
                solver.add(cs[i] == 0)
            else:
                solver.add(cs[i] != 0)

    and_mask, and_result, or_mask, or_result = 0, 0, 0, 0
    for i, (op, result) in enumerate(zip(ts, map(lambda v: (v >> bit) & 1, vs))):
        if op == 0:
            and_mask |= 1 << i
            if result == 1:
                and_result |= 1 << i
        else:
            or_mask |= 1 << i
            if result == 1:
                or_result |= 1 << i

    solver.add(reduce(and_, cs) & and_mask == and_result)
    solver.add(reduce(or_, cs) & or_mask == or_result)

    print("Solving bit {}...".format(bit), file=sys.stderr)
    solver.check()

※入出力を省略しています。完全なファイルはこちら

• 公式の問題解説にあるとおり、この問題はビットごとに独立して考えられるため、ビットごとに計算するよう修正
• 変数を$N\times N$個の1bit変数で持つのではなく、$N$個の$N$bit変数で持つよう修正

変数の消去など考察が必要な本質的な修正は行っていません。これにより、ギリギリ現実的な時間で終わるようにできました。

AtCoderに提出して確かめたいところですが、AtCoderのジャッジ環境でZ3は使えないので AtCoder's Testcases で公開されている実際の入出力を用いて手元で動かしました。

実行結果

anti-beet1.txt: 304.41s (AC)
anti-beet2.txt: 306.18s (AC)
anti-beet3.txt: 421.53s (AC)
anti-beet4.txt: 508.56s (AC)
anti-beet5.txt: 477.29s (AC)
anti-beet6.txt: 268.00s (AC)
anti-beet7.txt: 328.69s (AC)
anti-beet8.txt: 283.11s (AC)
anti_beet_alpha.txt: 261.81s (AC)
anti_beet_beta.txt: 295.70s (AC)
anti_beet_delta.txt: 304.54s (AC)
anti_beet_theta.txt: 316.23s (AC)
brute1.txt: 3.90s (AC)
brute2.txt: 3.87s (AC)
brute3.txt: 4.55s (AC)
brute4.txt: 4.10s (AC)
brute5.txt: 4.17s (AC)
brute6.txt: 4.96s (AC)
brute7.txt: 5.32s (AC)
brute8.txt: 4.00s (AC)
brute9.txt: 3.59s (AC)
ng3.txt: 199.05s (AC)
random1.txt: 397.36s (AC)
random2.txt: 323.34s (AC)
sample1.txt: 4.25s (AC)
sample2.txt: 3.75s (AC)
test0.txt: 1.75s (AC)
test10.txt: 404.20s (AC)
test11_.txt: 408.84s (AC)
test12_.txt: 384.60s (AC)
test13_.txt: 369.08s (AC)
test14.txt: 3.60s (AC)
test15.txt: 440.95s (AC)
test16.txt: 478.20s (AC)
test17_.txt: 410.48s (AC)
test18_.txt: 433.37s (AC)
test19_.txt: 439.01s (AC)
test1.txt: 18.18s (AC)
test20_.txt: 435.07s (AC)
test21_.txt: 417.01s (AC)
test22_.txt: 392.37s (AC)
test2.txt: 15.86s (AC)
test3.txt: 43.21s (AC)
test4.txt: 412.84s (AC)
test5.txt: 389.00s (AC)
test6.txt: 413.33s (AC)
test7.txt: 404.74s (AC)
test8.txt: 401.54s (AC)
test9.txt: 404.24s (AC)

一番遅いテストケースは508560msとなり、AtCoderの2000msの実行時間制限には遠く及びませんが、実際に問題の考察をすることなく正しい解を得られました。

SATソルバを用いる

SMTソルバは現実の計算に対して高度に抽象化された表現を受理できる一方、その代償として内部で無駄な計算を行っている可能性があります。

問題の「$i$行/列目の{論理和,論理積}が{1,0}である」という条件から、ビットごとの変数$x_i$を考えた時にこの問題に登場する制約は以下の4つです。

1. $\bigvee x_{i}=\top$
2. $\bigvee x_{i}=\bot$
3. $\bigwedge x_{i}=\top$
4. $\bigwedge x_{i}=\bot$

それぞれの制約は、以下のように直ちに連言標準形 (CNF) に変換できます。

1. $\bigvee x_{i}$
2. $\bigwedge \bar{x_{i}}$
3. $\bigwedge x_{i}$
4. $\bigvee \bar{x_{i}}$

すなわち、これら全ての制約を満たす論理式もそのままCNFとして解釈できます。難しい変換をするまでもなくそのままSATソルバに投げられる形になるので、SMTソルバではなく生のSATソルバを使って同じ実装をしてみましょう。SATソルバGlucose3を使いました。

ans = [0] * (n * n)

for bit in range(64):
    solver = Glucose3()

    for direction, (ops, results) in enumerate([[ss, us], [ts, vs]]):
        for i, (op, result) in enumerate(zip(ops, results)):
            if direction == 0:
                vals = [i * n + j + 1 for j in range(n)]
            else:
                vals = [j * n + i + 1 for j in range(n)]

            if op == 0:
                if (result >> bit) & 1 == 0:
                    solver.add_clause(list(map(neg, vals)))
                else:
                    for val in vals:
                        solver.add_clause([val])
            else:
                if (result >> bit) & 1 == 0:
                    for val in vals:
                        solver.add_clause([-val])
                else:
                    solver.add_clause(vals)

    print("Solving bit {}...".format(bit), file=sys.stderr)
    solver.solve()

※入出力を省略しています。完全なファイルはこちら

簡単ですね! 実際に動かすと、上のSMTソルバの10倍程度の速度で動作します。

実行結果

anti-beet1.txt: 40.84s (AC)
anti-beet2.txt: 38.06s (AC)
anti-beet3.txt: 36.02s (AC)
anti-beet4.txt: 35.78s (AC)
anti-beet5.txt: 38.89s (AC)
anti-beet6.txt: 46.40s (AC)
anti-beet7.txt: 48.26s (AC)
anti-beet8.txt: 40.96s (AC)
anti_beet_alpha.txt: 45.11s (AC)
anti_beet_beta.txt: 39.56s (AC)
anti_beet_delta.txt: 39.28s (AC)
anti_beet_theta.txt: 41.96s (AC)
brute1.txt: 1.77s (AC)
brute2.txt: 1.64s (AC)
brute3.txt: 1.66s (AC)
brute4.txt: 1.84s (AC)
brute5.txt: 1.22s (AC)
brute6.txt: 1.63s (AC)
brute7.txt: 1.25s (AC)
brute8.txt: 1.38s (AC)
brute9.txt: 1.30s (AC)
ng3.txt: 19.23s (AC)
random1.txt: 35.84s (AC)
random2.txt: 38.46s (AC)
sample1.txt: 1.25s (AC)
sample2.txt: 1.24s (AC)
small.txt: 1.34s (AC)
test0.txt: 1.63s (AC)
test10.txt: 35.73s (AC)
test11_.txt: 32.92s (AC)
test12_.txt: 34.10s (AC)
test13_.txt: 35.55s (AC)
test14.txt: 1.77s (AC)
test15.txt: 38.35s (AC)
test16.txt: 33.61s (AC)
test17_.txt: 34.22s (AC)
test18_.txt: 33.56s (AC)
test19_.txt: 34.98s (AC)
test1.txt: 3.39s (AC)
test20_.txt: 37.28s (AC)
test21_.txt: 32.69s (AC)
test22_.txt: 37.90s (AC)
test2.txt: 3.60s (AC)
test3.txt: 7.95s (AC)
test4.txt: 37.60s (AC)
test5.txt: 37.35s (AC)
test6.txt: 32.64s (AC)
test7.txt: 37.59s (AC)
test8.txt: 34.10s (AC)
test9.txt: 33.09s (AC)

一番遅いケースは48260msと、依然としてコンテストの制約には及びませんが、たかだか定数倍なので問題ありません。

ちなみにPySATから利用可能なSATソルバをすべて試しましたが、大きな性能の向上は見られませんでした。

• Glucose3: 38.337s
• Glucose4: 42.433s
• Cadical: 49.997s
• Lingeling: 54.917s
• MapleChrono: 43.778s
• MapleCM: 52.927s
• Maplesat: 44.063s
• Minicard: 37.287s
• Minisat22: 37.968s
• MinisatGH: 38.148s

まとめ

SMT/SATソルバはこのように満たしてほしい条件を入力するだけで全自動で問題を解いてくれる強力なツールです。自分のようにパズルが苦手な人種にはときおり重宝するツールなので、ぜひ一度触ってみてください。競プロには役に立たないと思いますが

なお競プロは詳しくないのでよくわかりませんが、2-SAT (連言標準形の節の長さがたかだか2であるような論理式の充足可能性問題) は実際に競技プログラミングで出題されることがあるようです。調べておくといいかもしれません。